


<!DOCTYPE html>
<!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]-->
<!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]-->
<head>
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>torch.utils.cpp_extension &mdash; PyTorch master documentation</title>
  

  
  
  
  
    <link rel="canonical" href="https://pytorch.org/docs/stable/_modules/torch/utils/cpp_extension.html"/>
  

  

  
  
    

  

  <link rel="stylesheet" href="../../../_static/css/theme.css" type="text/css" />
  <!-- <link rel="stylesheet" href="../../../_static/pygments.css" type="text/css" /> -->
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.10.0-beta/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/css/jit.css" type="text/css" />
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.11.1/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/katex-math.css" type="text/css" />
    <link rel="index" title="Index" href="../../../genindex.html" />
    <link rel="search" title="Search" href="../../../search.html" /> 

  
  <script src="../../../_static/js/modernizr.min.js"></script>

  <!-- Preload the theme fonts -->

<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-book.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/IBMPlexMono/IBMPlexMono-Medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-medium-italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/IBMPlexMono/IBMPlexMono-SemiBold.woff2" as="font" type="font/woff2" crossorigin="anonymous">

<!-- Preload the katex fonts -->

<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Math-Italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size1-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size4-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size2-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size3-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Caligraphic-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
</head>

<div class="container-fluid header-holder tutorials-header" id="header-holder">
  <div class="container">
    <div class="header-container">
      <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>

      <div class="main-menu">
        <ul>
          <li>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

          <li>
            <div class="ecosystem-dropdown">
              <a id="dropdownMenuButton" data-toggle="ecosystem-dropdown">
                Ecosystem
              </a>
              <div class="ecosystem-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/hub"">
                  <span class=dropdown-title>Models (Beta)</span>
                  <p>Discover, publish, and reuse pre-trained models</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/ecosystem">
                  <span class=dropdown-title>Tools & Libraries</span>
                  <p>Explore the ecosystem of tools and libraries</p>
                </a>
              </div>
            </div>
          </li>

          <li>
            <a href="https://pytorch.org/mobile">Mobile</a>
          </li>

          <li>
            <a href="https://pytorch.org/blog/">Blog</a>
          </li>

          <li>
            <a href="https://pytorch.org/tutorials">Tutorials</a>
          </li>

          <li class="active">
            <a href="https://pytorch.org/docs/stable/index.html">Docs</a>
          </li>

          <li>
            <div class="resources-dropdown">
              <a id="resourcesDropdownButton" data-toggle="resources-dropdown">
                Resources
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/resources"">
                  <span class=dropdown-title>Developer Resources</span>
                  <p>Find resources and get questions answered</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/features">
                  <span class=dropdown-title>About</span>
                  <p>Learn about PyTorch’s features and capabilities</p>
                </a>
              </div>
            </div>
          </li>

          <li>
            <a href="https://github.com/pytorch/pytorch">Github</a>
          </li>
        </ul>
      </div>

      <a class="main-menu-open-button" href="#" data-behavior="open-mobile-menu"></a>
    </div>

  </div>
</div>


<body class="pytorch-body">

   

    

    <div class="table-of-contents-link-wrapper">
      <span>Table of Contents</span>
      <a href="#" class="toggle-table-of-contents" data-behavior="toggle-table-of-contents"></a>
    </div>

    <nav data-toggle="wy-nav-shift" class="pytorch-left-menu" id="pytorch-left-menu">
      <div class="pytorch-side-scroll">
        <div class="pytorch-menu pytorch-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          <div class="pytorch-left-menu-search">
            

            
              
              
                <div class="version">
                  master (1.5.0 )
                </div>
              
            

            


  


<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search Docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

            
          </div>

          
<div>
  <a style="color:#F05732" href="https://pytorch.org/docs/stable/_modules/torch/utils/cpp_extension.html">
    You are viewing unstable developer preview docs.
    Click here to view docs for latest stable release.
  </a>
</div>

            
            
              
            
            
              <p class="caption"><span class="caption-text">Notes</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/amp_examples.html">Automatic Mixed Precision examples</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/autograd.html">Autograd mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/broadcasting.html">Broadcasting semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/cpu_threading_torchscript_inference.html">CPU threading and TorchScript inference</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/cuda.html">CUDA semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/ddp.html">Distributed Data Parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/extending.html">Extending PyTorch</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/faq.html">Frequently Asked Questions</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/large_scale_deployments.html">Features for large-scale deployments</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/multiprocessing.html">Multiprocessing best practices</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/randomness.html">Reproducibility</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/serialization.html">Serialization semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/windows.html">Windows FAQ</a></li>
</ul>
<p class="caption"><span class="caption-text">Language Bindings</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/cppdocs/">C++ API</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../packages.html">Javadoc</a></li>
</ul>
<p class="caption"><span class="caption-text">Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../torch.html">torch</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../nn.html">torch.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../nn.functional.html">torch.nn.functional</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensors.html">torch.Tensor</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensor_attributes.html">Tensor Attributes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensor_view.html">Tensor Views</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../autograd.html">torch.autograd</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../cuda.html">torch.cuda</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../amp.html">torch.cuda.amp</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributed.html">torch.distributed</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributions.html">torch.distributions</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../hub.html">torch.hub</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../jit.html">torch.jit</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../nn.init.html">torch.nn.init</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../onnx.html">torch.onnx</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../optim.html">torch.optim</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../quantization.html">Quantization</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../rpc/index.html">Distributed RPC Framework</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../random.html">torch.random</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../sparse.html">torch.sparse</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../storage.html">torch.Storage</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../bottleneck.html">torch.utils.bottleneck</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../checkpoint.html">torch.utils.checkpoint</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../cpp_extension.html">torch.utils.cpp_extension</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../data.html">torch.utils.data</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../dlpack.html">torch.utils.dlpack</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../model_zoo.html">torch.utils.model_zoo</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensorboard.html">torch.utils.tensorboard</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../type_info.html">Type Info</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../named_tensor.html">Named Tensors</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../name_inference.html">Named Tensors operator coverage</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../__config__.html">torch.__config__</a></li>
</ul>
<p class="caption"><span class="caption-text">Libraries</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/audio">torchaudio</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/text">torchtext</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/elastic/">TorchElastic</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/serve">TorchServe</a></li>
<li class="toctree-l1"><a class="reference external" href="http://pytorch.org/xla/">PyTorch on XLA Devices</a></li>
</ul>
<p class="caption"><span class="caption-text">Community</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../community/contribution_guide.html">PyTorch Contribution Guide</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../community/governance.html">PyTorch Governance</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../community/persons_of_interest.html">PyTorch Governance | Persons of Interest</a></li>
</ul>

            
          

        </div>
      </div>
    </nav>

    <div class="pytorch-container">
      <div class="pytorch-page-level-bar" id="pytorch-page-level-bar">
        <div class="pytorch-breadcrumbs-wrapper">
          















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="pytorch-breadcrumbs">
    
      <li>
        <a href="../../../index.html">
          
            Docs
          
        </a> &gt;
      </li>

        
          <li><a href="../../index.html">Module code</a> &gt;</li>
        
          <li><a href="../../torch.html">torch</a> &gt;</li>
        
      <li>torch.utils.cpp_extension</li>
    
    
      <li class="pytorch-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
</div>
        </div>

        <div class="pytorch-shortcuts-wrapper" id="pytorch-shortcuts-wrapper">
          Shortcuts
        </div>
      </div>

      <section data-toggle="wy-nav-shift" id="pytorch-content-wrap" class="pytorch-content-wrap">
        <div class="pytorch-content-left">

        
          
          <div class="rst-content">
          
            <div role="main" class="main-content" itemscope="itemscope" itemtype="http://schema.org/Article">
             <article itemprop="articleBody" id="pytorch-article" class="pytorch-article">
              
  <h1>Source code for torch.utils.cpp_extension</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">__future__</span> <span class="kn">import</span> <span class="n">absolute_import</span><span class="p">,</span> <span class="n">division</span><span class="p">,</span> <span class="n">print_function</span><span class="p">,</span> <span class="n">unicode_literals</span>
<span class="kn">import</span> <span class="nn">copy</span>
<span class="kn">import</span> <span class="nn">glob</span>
<span class="kn">import</span> <span class="nn">imp</span>
<span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">re</span>
<span class="kn">import</span> <span class="nn">setuptools</span>
<span class="kn">import</span> <span class="nn">subprocess</span>
<span class="kn">import</span> <span class="nn">sys</span>
<span class="kn">import</span> <span class="nn">sysconfig</span>
<span class="kn">import</span> <span class="nn">tempfile</span>
<span class="kn">import</span> <span class="nn">warnings</span>
<span class="kn">import</span> <span class="nn">collections</span>

<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">.file_baton</span> <span class="kn">import</span> <span class="n">FileBaton</span>
<span class="kn">from</span> <span class="nn">._cpp_extension_versioner</span> <span class="kn">import</span> <span class="n">ExtensionVersioner</span>

<span class="kn">from</span> <span class="nn">setuptools.command.build_ext</span> <span class="kn">import</span> <span class="n">build_ext</span>


<span class="n">IS_WINDOWS</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">platform</span> <span class="o">==</span> <span class="s1">&#39;win32&#39;</span>

<span class="k">def</span> <span class="nf">_find_cuda_home</span><span class="p">():</span>
    <span class="sd">&#39;&#39;&#39;Finds the CUDA install path.&#39;&#39;&#39;</span>
    <span class="c1"># Guess #1</span>
    <span class="n">cuda_home</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CUDA_HOME&#39;</span><span class="p">)</span> <span class="ow">or</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CUDA_PATH&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">cuda_home</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="c1"># Guess #2</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">which</span> <span class="o">=</span> <span class="s1">&#39;where&#39;</span> <span class="k">if</span> <span class="n">IS_WINDOWS</span> <span class="k">else</span> <span class="s1">&#39;which&#39;</span>
            <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">devnull</span><span class="p">,</span> <span class="s1">&#39;w&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">devnull</span><span class="p">:</span>
                <span class="n">nvcc</span> <span class="o">=</span> <span class="n">subprocess</span><span class="o">.</span><span class="n">check_output</span><span class="p">([</span><span class="n">which</span><span class="p">,</span> <span class="s1">&#39;nvcc&#39;</span><span class="p">],</span>
                                               <span class="n">stderr</span><span class="o">=</span><span class="n">devnull</span><span class="p">)</span><span class="o">.</span><span class="n">decode</span><span class="p">()</span><span class="o">.</span><span class="n">rstrip</span><span class="p">(</span><span class="s1">&#39;</span><span class="se">\r\n</span><span class="s1">&#39;</span><span class="p">)</span>
                <span class="n">cuda_home</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">nvcc</span><span class="p">))</span>
        <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
            <span class="c1"># Guess #3</span>
            <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
                <span class="n">cuda_homes</span> <span class="o">=</span> <span class="n">glob</span><span class="o">.</span><span class="n">glob</span><span class="p">(</span>
                    <span class="s1">&#39;C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v*.*&#39;</span><span class="p">)</span>
                <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">cuda_homes</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">cuda_home</span> <span class="o">=</span> <span class="s1">&#39;&#39;</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">cuda_home</span> <span class="o">=</span> <span class="n">cuda_homes</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">cuda_home</span> <span class="o">=</span> <span class="s1">&#39;/usr/local/cuda&#39;</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">cuda_home</span><span class="p">):</span>
                <span class="n">cuda_home</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">if</span> <span class="n">cuda_home</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">is_available</span><span class="p">():</span>
        <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;No CUDA runtime is found, using CUDA_HOME=&#39;</span><span class="si">{}</span><span class="s2">&#39;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">cuda_home</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">cuda_home</span>

<span class="k">def</span> <span class="nf">_find_rocm_home</span><span class="p">():</span>
    <span class="sd">&#39;&#39;&#39;Finds the ROCm install path.&#39;&#39;&#39;</span>
    <span class="c1"># Guess #1</span>
    <span class="n">rocm_home</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;ROCM_HOME&#39;</span><span class="p">)</span> <span class="ow">or</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;ROCM_PATH&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">rocm_home</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="c1"># Guess #2</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">hipcc</span> <span class="o">=</span> <span class="n">subprocess</span><span class="o">.</span><span class="n">check_output</span><span class="p">(</span>
                <span class="p">[</span><span class="s1">&#39;which&#39;</span><span class="p">,</span> <span class="s1">&#39;hipcc&#39;</span><span class="p">])</span><span class="o">.</span><span class="n">decode</span><span class="p">()</span><span class="o">.</span><span class="n">rstrip</span><span class="p">(</span><span class="s1">&#39;</span><span class="se">\r\n</span><span class="s1">&#39;</span><span class="p">)</span>
            <span class="c1"># this will be either &lt;ROCM_HOME&gt;/hip/bin/hipcc or &lt;ROCM_HOME&gt;/bin/hipcc</span>
            <span class="n">rocm_home</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">hipcc</span><span class="p">))</span>
            <span class="k">if</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">basename</span><span class="p">(</span><span class="n">rocm_home</span><span class="p">)</span> <span class="o">==</span> <span class="s1">&#39;hip&#39;</span><span class="p">:</span>
                <span class="n">rocm_home</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">rocm_home</span><span class="p">)</span>
        <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
            <span class="c1"># Guess #3</span>
            <span class="n">rocm_home</span> <span class="o">=</span> <span class="s1">&#39;/opt/rocm&#39;</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">rocm_home</span><span class="p">):</span>
                <span class="n">rocm_home</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">if</span> <span class="n">rocm_home</span> <span class="ow">and</span> <span class="n">torch</span><span class="o">.</span><span class="n">version</span><span class="o">.</span><span class="n">hip</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;No ROCm runtime is found, using ROCM_HOME=&#39;</span><span class="si">{}</span><span class="s2">&#39;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">rocm_home</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">rocm_home</span>


<span class="k">def</span> <span class="nf">_join_rocm_home</span><span class="p">(</span><span class="o">*</span><span class="n">paths</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Joins paths with ROCM_HOME, or raises an error if it ROCM_HOME is not set.</span>

<span class="sd">    This is basically a lazy way of raising an error for missing $ROCM_HOME</span>
<span class="sd">    only once we need to get any ROCm-specific path.</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="k">if</span> <span class="n">ROCM_HOME</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">EnvironmentError</span><span class="p">(</span><span class="s1">&#39;ROCM_HOME environment variable is not set. &#39;</span>
                               <span class="s1">&#39;Please set it to your ROCm install root.&#39;</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">EnvironmentError</span><span class="p">(</span><span class="s1">&#39;Building PyTorch extensions using &#39;</span>
                               <span class="s1">&#39;ROCm and Windows is not supported.&#39;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">ROCM_HOME</span><span class="p">,</span> <span class="o">*</span><span class="n">paths</span><span class="p">)</span>


<span class="n">MINIMUM_GCC_VERSION</span> <span class="o">=</span> <span class="p">(</span><span class="mi">5</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
<span class="n">MINIMUM_MSVC_VERSION</span> <span class="o">=</span> <span class="p">(</span><span class="mi">19</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">24215</span><span class="p">)</span>
<span class="n">ABI_INCOMPATIBILITY_WARNING</span> <span class="o">=</span> <span class="s1">&#39;&#39;&#39;</span>

<span class="s1">                               !! WARNING !!</span>

<span class="s1">!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!</span>
<span class="s1">Your compiler (</span><span class="si">{}</span><span class="s1">) may be ABI-incompatible with PyTorch!</span>
<span class="s1">Please use a compiler that is ABI-compatible with GCC 5.0 and above.</span>
<span class="s1">See https://gcc.gnu.org/onlinedocs/libstdc++/manual/abi.html.</span>

<span class="s1">See https://gist.github.com/goldsborough/d466f43e8ffc948ff92de7486c5216d6</span>
<span class="s1">for instructions on how to install GCC 5 or higher.</span>
<span class="s1">!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!</span>

<span class="s1">                              !! WARNING !!</span>
<span class="s1">&#39;&#39;&#39;</span>
<span class="n">WRONG_COMPILER_WARNING</span> <span class="o">=</span> <span class="s1">&#39;&#39;&#39;</span>

<span class="s1">                               !! WARNING !!</span>

<span class="s1">!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!</span>
<span class="s1">Your compiler (</span><span class="si">{user_compiler}</span><span class="s1">) is not compatible with the compiler Pytorch was</span>
<span class="s1">built with for this platform, which is </span><span class="si">{pytorch_compiler}</span><span class="s1"> on </span><span class="si">{platform}</span><span class="s1">. Please</span>
<span class="s1">use </span><span class="si">{pytorch_compiler}</span><span class="s1"> to to compile your extension. Alternatively, you may</span>
<span class="s1">compile PyTorch from source using </span><span class="si">{user_compiler}</span><span class="s1">, and then you can also use</span>
<span class="si">{user_compiler}</span><span class="s1"> to compile your extension.</span>

<span class="s1">See https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md for help</span>
<span class="s1">with compiling PyTorch from source.</span>
<span class="s1">!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!</span>

<span class="s1">                              !! WARNING !!</span>
<span class="s1">&#39;&#39;&#39;</span>
<span class="n">ROCM_HOME</span> <span class="o">=</span> <span class="n">_find_rocm_home</span><span class="p">()</span>
<span class="n">MIOPEN_HOME</span> <span class="o">=</span> <span class="n">_join_rocm_home</span><span class="p">(</span><span class="s1">&#39;miopen&#39;</span><span class="p">)</span> <span class="k">if</span> <span class="n">ROCM_HOME</span> <span class="k">else</span> <span class="kc">None</span>
<span class="n">IS_HIP_EXTENSION</span> <span class="o">=</span> <span class="kc">True</span> <span class="k">if</span> <span class="p">((</span><span class="n">ROCM_HOME</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">)</span> <span class="ow">and</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">version</span><span class="o">.</span><span class="n">hip</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">))</span> <span class="k">else</span> <span class="kc">False</span>
<span class="n">CUDA_HOME</span> <span class="o">=</span> <span class="n">_find_cuda_home</span><span class="p">()</span>
<span class="n">CUDNN_HOME</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CUDNN_HOME&#39;</span><span class="p">)</span> <span class="ow">or</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CUDNN_PATH&#39;</span><span class="p">)</span>
<span class="c1"># PyTorch releases have the version pattern major.minor.patch, whereas when</span>
<span class="c1"># PyTorch is built from source, we append the git commit hash, which gives</span>
<span class="c1"># it the below pattern.</span>
<span class="n">BUILT_FROM_SOURCE_VERSION_PATTERN</span> <span class="o">=</span> <span class="n">re</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="sa">r</span><span class="s1">&#39;\d+\.\d+\.\d+\w+\+\w+&#39;</span><span class="p">)</span>

<span class="n">COMMON_MSVC_FLAGS</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;/MD&#39;</span><span class="p">,</span> <span class="s1">&#39;/wd4819&#39;</span><span class="p">,</span> <span class="s1">&#39;/EHsc&#39;</span><span class="p">]</span>

<span class="n">COMMON_NVCC_FLAGS</span> <span class="o">=</span> <span class="p">[</span>
    <span class="s1">&#39;-D__CUDA_NO_HALF_OPERATORS__&#39;</span><span class="p">,</span>
    <span class="s1">&#39;-D__CUDA_NO_HALF_CONVERSIONS__&#39;</span><span class="p">,</span>
    <span class="s1">&#39;-D__CUDA_NO_HALF2_OPERATORS__&#39;</span><span class="p">,</span>
    <span class="s1">&#39;--expt-relaxed-constexpr&#39;</span>
<span class="p">]</span>

<span class="n">COMMON_HIPCC_FLAGS</span> <span class="o">=</span> <span class="p">[</span>
    <span class="s1">&#39;-fPIC&#39;</span><span class="p">,</span>
    <span class="s1">&#39;-D__HIP_PLATFORM_HCC__=1&#39;</span><span class="p">,</span>
    <span class="s1">&#39;-DCUDA_HAS_FP16=1&#39;</span><span class="p">,</span>
    <span class="s1">&#39;-D__HIP_NO_HALF_OPERATORS__=1&#39;</span><span class="p">,</span>
    <span class="s1">&#39;-D__HIP_NO_HALF_CONVERSIONS__=1&#39;</span><span class="p">,</span>
<span class="p">]</span>

<span class="n">JIT_EXTENSION_VERSIONER</span> <span class="o">=</span> <span class="n">ExtensionVersioner</span><span class="p">()</span>


<span class="k">def</span> <span class="nf">_is_binary_build</span><span class="p">():</span>
    <span class="k">return</span> <span class="ow">not</span> <span class="n">BUILT_FROM_SOURCE_VERSION_PATTERN</span><span class="o">.</span><span class="n">match</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">version</span><span class="o">.</span><span class="n">__version__</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_accepted_compilers_for_platform</span><span class="p">():</span>
    <span class="c1"># gnu-c++ and gnu-cc are the conda gcc compilers</span>
    <span class="k">return</span> <span class="p">[</span><span class="s1">&#39;clang++&#39;</span><span class="p">,</span> <span class="s1">&#39;clang&#39;</span><span class="p">]</span> <span class="k">if</span> <span class="n">sys</span><span class="o">.</span><span class="n">platform</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s1">&#39;darwin&#39;</span><span class="p">)</span> <span class="k">else</span> <span class="p">[</span><span class="s1">&#39;g++&#39;</span><span class="p">,</span> <span class="s1">&#39;gcc&#39;</span><span class="p">,</span> <span class="s1">&#39;gnu-c++&#39;</span><span class="p">,</span> <span class="s1">&#39;gnu-cc&#39;</span><span class="p">]</span>


<span class="k">def</span> <span class="nf">get_default_build_root</span><span class="p">():</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Returns the path to the root folder under which extensions will built.</span>

<span class="sd">    For each extension module built, there will be one folder underneath the</span>
<span class="sd">    folder returned by this function. For example, if ``p`` is the path</span>
<span class="sd">    returned by this function and ``ext`` the name of an extension, the build</span>
<span class="sd">    folder for the extension will be ``p/ext``.</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="c1"># tempfile.gettempdir() will be /tmp on UNIX and \TEMP on Windows.</span>
    <span class="k">return</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">realpath</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">tempfile</span><span class="o">.</span><span class="n">gettempdir</span><span class="p">(),</span> <span class="s1">&#39;torch_extensions&#39;</span><span class="p">))</span>


<span class="k">def</span> <span class="nf">check_compiler_ok_for_platform</span><span class="p">(</span><span class="n">compiler</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Verifies that the compiler is the expected one for the current platform.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        compiler (str): The compiler executable to check.</span>

<span class="sd">    Returns:</span>
<span class="sd">        True if the compiler is gcc/g++ on Linux or clang/clang++ on macOS,</span>
<span class="sd">        and always True for Windows.</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
        <span class="k">return</span> <span class="kc">True</span>
    <span class="n">which</span> <span class="o">=</span> <span class="n">subprocess</span><span class="o">.</span><span class="n">check_output</span><span class="p">([</span><span class="s1">&#39;which&#39;</span><span class="p">,</span> <span class="n">compiler</span><span class="p">],</span> <span class="n">stderr</span><span class="o">=</span><span class="n">subprocess</span><span class="o">.</span><span class="n">STDOUT</span><span class="p">)</span>
    <span class="c1"># Use os.path.realpath to resolve any symlinks, in particular from &#39;c++&#39; to e.g. &#39;g++&#39;.</span>
    <span class="n">compiler_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">realpath</span><span class="p">(</span><span class="n">which</span><span class="o">.</span><span class="n">decode</span><span class="p">()</span><span class="o">.</span><span class="n">strip</span><span class="p">())</span>
    <span class="k">return</span> <span class="nb">any</span><span class="p">(</span><span class="n">name</span> <span class="ow">in</span> <span class="n">compiler_path</span> <span class="k">for</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">_accepted_compilers_for_platform</span><span class="p">())</span>


<div class="viewcode-block" id="check_compiler_abi_compatibility"><a class="viewcode-back" href="../../../cpp_extension.html#torch.utils.cpp_extension.check_compiler_abi_compatibility">[docs]</a><span class="k">def</span> <span class="nf">check_compiler_abi_compatibility</span><span class="p">(</span><span class="n">compiler</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Verifies that the given compiler is ABI-compatible with PyTorch.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        compiler (str): The compiler executable name to check (e.g. ``g++``).</span>
<span class="sd">            Must be executable in a shell process.</span>

<span class="sd">    Returns:</span>
<span class="sd">        False if the compiler is (likely) ABI-incompatible with PyTorch,</span>
<span class="sd">        else True.</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">_is_binary_build</span><span class="p">():</span>
        <span class="k">return</span> <span class="kc">True</span>
    <span class="k">if</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;TORCH_DONT_CHECK_COMPILER_ABI&#39;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="s1">&#39;ON&#39;</span><span class="p">,</span> <span class="s1">&#39;1&#39;</span><span class="p">,</span> <span class="s1">&#39;YES&#39;</span><span class="p">,</span> <span class="s1">&#39;TRUE&#39;</span><span class="p">,</span> <span class="s1">&#39;Y&#39;</span><span class="p">]:</span>
        <span class="k">return</span> <span class="kc">True</span>

    <span class="c1"># First check if the compiler is one of the expected ones for the particular platform.</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">check_compiler_ok_for_platform</span><span class="p">(</span><span class="n">compiler</span><span class="p">):</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="n">WRONG_COMPILER_WARNING</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
            <span class="n">user_compiler</span><span class="o">=</span><span class="n">compiler</span><span class="p">,</span>
            <span class="n">pytorch_compiler</span><span class="o">=</span><span class="n">_accepted_compilers_for_platform</span><span class="p">()[</span><span class="mi">0</span><span class="p">],</span>
            <span class="n">platform</span><span class="o">=</span><span class="n">sys</span><span class="o">.</span><span class="n">platform</span><span class="p">))</span>
        <span class="k">return</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="n">sys</span><span class="o">.</span><span class="n">platform</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s1">&#39;darwin&#39;</span><span class="p">):</span>
        <span class="c1"># There is no particular minimum version we need for clang, so we&#39;re good here.</span>
        <span class="k">return</span> <span class="kc">True</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">sys</span><span class="o">.</span><span class="n">platform</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s1">&#39;linux&#39;</span><span class="p">):</span>
            <span class="n">minimum_required_version</span> <span class="o">=</span> <span class="n">MINIMUM_GCC_VERSION</span>
            <span class="n">version</span> <span class="o">=</span> <span class="n">subprocess</span><span class="o">.</span><span class="n">check_output</span><span class="p">([</span><span class="n">compiler</span><span class="p">,</span> <span class="s1">&#39;-dumpfullversion&#39;</span><span class="p">,</span> <span class="s1">&#39;-dumpversion&#39;</span><span class="p">])</span>
            <span class="n">version</span> <span class="o">=</span> <span class="n">version</span><span class="o">.</span><span class="n">decode</span><span class="p">()</span><span class="o">.</span><span class="n">strip</span><span class="p">()</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s1">&#39;.&#39;</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">minimum_required_version</span> <span class="o">=</span> <span class="n">MINIMUM_MSVC_VERSION</span>
            <span class="n">compiler_info</span> <span class="o">=</span> <span class="n">subprocess</span><span class="o">.</span><span class="n">check_output</span><span class="p">(</span><span class="n">compiler</span><span class="p">,</span> <span class="n">stderr</span><span class="o">=</span><span class="n">subprocess</span><span class="o">.</span><span class="n">STDOUT</span><span class="p">)</span>
            <span class="n">match</span> <span class="o">=</span> <span class="n">re</span><span class="o">.</span><span class="n">search</span><span class="p">(</span><span class="sa">r</span><span class="s1">&#39;(\d+)\.(\d+)\.(\d+)&#39;</span><span class="p">,</span> <span class="n">compiler_info</span><span class="o">.</span><span class="n">decode</span><span class="p">()</span><span class="o">.</span><span class="n">strip</span><span class="p">())</span>
            <span class="n">version</span> <span class="o">=</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span> <span class="k">if</span> <span class="n">match</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">match</span><span class="o">.</span><span class="n">groups</span><span class="p">()</span>
    <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
        <span class="n">_</span><span class="p">,</span> <span class="n">error</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">exc_info</span><span class="p">()</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s1">&#39;Error checking compiler version for </span><span class="si">{}</span><span class="s1">: </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">compiler</span><span class="p">,</span> <span class="n">error</span><span class="p">))</span>
        <span class="k">return</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="nb">int</span><span class="p">,</span> <span class="n">version</span><span class="p">))</span> <span class="o">&gt;=</span> <span class="n">minimum_required_version</span><span class="p">:</span>
        <span class="k">return</span> <span class="kc">True</span>

    <span class="n">compiler</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="si">{}</span><span class="s1"> </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">compiler</span><span class="p">,</span> <span class="s2">&quot;.&quot;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">version</span><span class="p">))</span>
    <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="n">ABI_INCOMPATIBILITY_WARNING</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">compiler</span><span class="p">))</span>

    <span class="k">return</span> <span class="kc">False</span></div>


<span class="c1"># See below for why we inherit BuildExtension from object.</span>
<span class="c1"># https://stackoverflow.com/questions/1713038/super-fails-with-error-typeerror-argument-1-must-be-type-not-classobj-when</span>


<div class="viewcode-block" id="BuildExtension"><a class="viewcode-back" href="../../../cpp_extension.html#torch.utils.cpp_extension.BuildExtension">[docs]</a><span class="k">class</span> <span class="nc">BuildExtension</span><span class="p">(</span><span class="n">build_ext</span><span class="p">,</span> <span class="nb">object</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    A custom :mod:`setuptools` build extension .</span>

<span class="sd">    This :class:`setuptools.build_ext` subclass takes care of passing the</span>
<span class="sd">    minimum required compiler flags (e.g. ``-std=c++14``) as well as mixed</span>
<span class="sd">    C++/CUDA compilation (and support for CUDA files in general).</span>

<span class="sd">    When using :class:`BuildExtension`, it is allowed to supply a dictionary</span>
<span class="sd">    for ``extra_compile_args`` (rather than the usual list) that maps from</span>
<span class="sd">    languages (``cxx`` or ``nvcc``) to a list of additional compiler flags to</span>
<span class="sd">    supply to the compiler. This makes it possible to supply different flags to</span>
<span class="sd">    the C++ and CUDA compiler during mixed compilation.</span>

<span class="sd">    ``use_ninja`` (bool): If ``use_ninja`` is ``True`` (default), then we</span>
<span class="sd">    attempt to build using the Ninja backend. Ninja greatly speeds up</span>
<span class="sd">    compilation compared to the standard ``setuptools.build_ext``.</span>
<span class="sd">    Fallbacks to the standard distutils backend if Ninja is not available.</span>

<span class="sd">    .. note::</span>
<span class="sd">        By default, the Ninja backend uses #CPUS + 2 workers to build the</span>
<span class="sd">        extension. This may use up too many resources on some systems. One</span>
<span class="sd">        can control the number of workers by setting the `MAX_JOBS` environment</span>
<span class="sd">        variable to a non-negative number.</span>
<span class="sd">    &#39;&#39;&#39;</span>

    <span class="nd">@classmethod</span>
    <span class="k">def</span> <span class="nf">with_options</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="o">**</span><span class="n">options</span><span class="p">):</span>
        <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">        Returns an alternative constructor that extends any original keyword</span>
<span class="sd">        arguments to the original constructor with the given options.</span>
<span class="sd">        &#39;&#39;&#39;</span>
        <span class="k">def</span> <span class="nf">init_with_options</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
            <span class="n">kwargs</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
            <span class="n">kwargs</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">options</span><span class="p">)</span>
            <span class="k">return</span> <span class="bp">cls</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">init_with_options</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">BuildExtension</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">no_python_abi_suffix</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;no_python_abi_suffix&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">use_ninja</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;use_ninja&#39;</span><span class="p">,</span> <span class="kc">False</span> <span class="k">if</span> <span class="n">IS_HIP_EXTENSION</span> <span class="k">else</span> <span class="kc">True</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">use_ninja</span><span class="p">:</span>
            <span class="c1"># Test if we can use ninja. Fallback otherwise.</span>
            <span class="n">msg</span> <span class="o">=</span> <span class="p">(</span><span class="s1">&#39;Attempted to use ninja as the BuildExtension backend but &#39;</span>
                   <span class="s1">&#39;</span><span class="si">{}</span><span class="s1">. Falling back to using the slow distutils backend.&#39;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">IS_HIP_EXTENSION</span><span class="p">:</span>
                <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="n">msg</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="s1">&#39;HIP extensions is not supported yet for ninja.&#39;</span><span class="p">))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">use_ninja</span> <span class="o">=</span> <span class="kc">False</span>
            <span class="k">elif</span> <span class="ow">not</span> <span class="n">_is_ninja_available</span><span class="p">():</span>
                <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="n">msg</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="s1">&#39;we could not find ninja.&#39;</span><span class="p">))</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">use_ninja</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">def</span> <span class="nf">build_extensions</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_check_abi</span><span class="p">()</span>
        <span class="k">for</span> <span class="n">extension</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">extensions</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_add_compile_flag</span><span class="p">(</span><span class="n">extension</span><span class="p">,</span> <span class="s1">&#39;-DTORCH_API_INCLUDE_EXTENSION_H&#39;</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_define_torch_extension_name</span><span class="p">(</span><span class="n">extension</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_add_gnu_cpp_abi_flag</span><span class="p">(</span><span class="n">extension</span><span class="p">)</span>

        <span class="c1"># Register .cu, .cuh and .hip as valid source extensions.</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">src_extensions</span> <span class="o">+=</span> <span class="p">[</span><span class="s1">&#39;.cu&#39;</span><span class="p">,</span> <span class="s1">&#39;.cuh&#39;</span><span class="p">,</span> <span class="s1">&#39;.hip&#39;</span><span class="p">]</span>
        <span class="c1"># Save the original _compile method for later.</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compiler_type</span> <span class="o">==</span> <span class="s1">&#39;msvc&#39;</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">_cpp_extensions</span> <span class="o">+=</span> <span class="p">[</span><span class="s1">&#39;.cu&#39;</span><span class="p">,</span> <span class="s1">&#39;.cuh&#39;</span><span class="p">]</span>
            <span class="n">original_compile</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compile</span>
            <span class="n">original_spawn</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">spawn</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">original_compile</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">_compile</span>

        <span class="k">def</span> <span class="nf">append_std14_if_no_std_present</span><span class="p">(</span><span class="n">cflags</span><span class="p">):</span>
            <span class="c1"># NVCC does not allow multiple -std to be passed, so we avoid</span>
            <span class="c1"># overriding the option if the user explicitly passed it.</span>
            <span class="n">cpp_format_prefix</span> <span class="o">=</span> <span class="s1">&#39;/</span><span class="si">{}</span><span class="s1">:&#39;</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compiler_type</span> <span class="o">==</span> <span class="s1">&#39;msvc&#39;</span> <span class="k">else</span> <span class="s1">&#39;-</span><span class="si">{}</span><span class="s1">=&#39;</span>
            <span class="n">cpp_flag_prefix</span> <span class="o">=</span> <span class="n">cpp_format_prefix</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="s1">&#39;std&#39;</span><span class="p">)</span>
            <span class="n">cpp_flag</span> <span class="o">=</span> <span class="n">cpp_flag_prefix</span> <span class="o">+</span> <span class="s1">&#39;c++14&#39;</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">flag</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="n">cpp_flag_prefix</span><span class="p">)</span> <span class="k">for</span> <span class="n">flag</span> <span class="ow">in</span> <span class="n">cflags</span><span class="p">):</span>
                <span class="n">cflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">cpp_flag</span><span class="p">)</span>

        <span class="k">def</span> <span class="nf">unix_cuda_flags</span><span class="p">(</span><span class="n">cflags</span><span class="p">):</span>
            <span class="k">return</span> <span class="p">(</span><span class="n">COMMON_NVCC_FLAGS</span> <span class="o">+</span>
                    <span class="p">[</span><span class="s1">&#39;--compiler-options&#39;</span><span class="p">,</span> <span class="s2">&quot;&#39;-fPIC&#39;&quot;</span><span class="p">]</span> <span class="o">+</span>
                    <span class="n">cflags</span> <span class="o">+</span> <span class="n">_get_cuda_arch_flags</span><span class="p">(</span><span class="n">cflags</span><span class="p">))</span>

        <span class="k">def</span> <span class="nf">unix_wrap_single_compile</span><span class="p">(</span><span class="n">obj</span><span class="p">,</span> <span class="n">src</span><span class="p">,</span> <span class="n">ext</span><span class="p">,</span> <span class="n">cc_args</span><span class="p">,</span> <span class="n">extra_postargs</span><span class="p">,</span> <span class="n">pp_opts</span><span class="p">):</span>
            <span class="c1"># Copy before we make any modifications.</span>
            <span class="n">cflags</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">extra_postargs</span><span class="p">)</span>
            <span class="k">try</span><span class="p">:</span>
                <span class="n">original_compiler</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compiler_so</span>
                <span class="k">if</span> <span class="n">_is_cuda_file</span><span class="p">(</span><span class="n">src</span><span class="p">):</span>
                    <span class="n">nvcc</span> <span class="o">=</span> <span class="p">(</span><span class="n">_join_rocm_home</span><span class="p">(</span><span class="s1">&#39;bin&#39;</span><span class="p">,</span> <span class="s1">&#39;hipcc&#39;</span><span class="p">)</span> <span class="k">if</span> <span class="n">IS_HIP_EXTENSION</span> <span class="k">else</span> <span class="n">_join_cuda_home</span><span class="p">(</span><span class="s1">&#39;bin&#39;</span><span class="p">,</span> <span class="s1">&#39;nvcc&#39;</span><span class="p">))</span>
                    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">nvcc</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
                        <span class="n">nvcc</span> <span class="o">=</span> <span class="p">[</span><span class="n">nvcc</span><span class="p">]</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">set_executable</span><span class="p">(</span><span class="s1">&#39;compiler_so&#39;</span><span class="p">,</span> <span class="n">nvcc</span><span class="p">)</span>
                    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">cflags</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
                        <span class="n">cflags</span> <span class="o">=</span> <span class="n">cflags</span><span class="p">[</span><span class="s1">&#39;nvcc&#39;</span><span class="p">]</span>
                    <span class="k">if</span> <span class="n">IS_HIP_EXTENSION</span><span class="p">:</span>
                        <span class="n">cflags</span> <span class="o">=</span> <span class="n">cflags</span> <span class="o">+</span> <span class="n">_get_rocm_arch_flags</span><span class="p">(</span><span class="n">cflags</span><span class="p">)</span>
                    <span class="k">else</span><span class="p">:</span>
                        <span class="n">cflags</span> <span class="o">=</span> <span class="n">unix_cuda_flags</span><span class="p">(</span><span class="n">cflags</span><span class="p">)</span>
                <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">cflags</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
                    <span class="n">cflags</span> <span class="o">=</span> <span class="n">cflags</span><span class="p">[</span><span class="s1">&#39;cxx&#39;</span><span class="p">]</span>
                <span class="k">if</span> <span class="n">IS_HIP_EXTENSION</span><span class="p">:</span>
                    <span class="n">cflags</span> <span class="o">=</span> <span class="n">cflags</span> <span class="o">+</span> <span class="n">COMMON_HIPCC_FLAGS</span>
                <span class="n">append_std14_if_no_std_present</span><span class="p">(</span><span class="n">cflags</span><span class="p">)</span>

                <span class="n">original_compile</span><span class="p">(</span><span class="n">obj</span><span class="p">,</span> <span class="n">src</span><span class="p">,</span> <span class="n">ext</span><span class="p">,</span> <span class="n">cc_args</span><span class="p">,</span> <span class="n">cflags</span><span class="p">,</span> <span class="n">pp_opts</span><span class="p">)</span>
            <span class="k">finally</span><span class="p">:</span>
                <span class="c1"># Put the original compiler back in place.</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">set_executable</span><span class="p">(</span><span class="s1">&#39;compiler_so&#39;</span><span class="p">,</span> <span class="n">original_compiler</span><span class="p">)</span>

        <span class="k">def</span> <span class="nf">unix_wrap_ninja_compile</span><span class="p">(</span><span class="n">sources</span><span class="p">,</span>
                                    <span class="n">output_dir</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                    <span class="n">macros</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                    <span class="n">include_dirs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                    <span class="n">debug</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                                    <span class="n">extra_preargs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                    <span class="n">extra_postargs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                    <span class="n">depends</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
            <span class="sd">&quot;&quot;&quot;Compiles sources by outputting a ninja file and running it.&quot;&quot;&quot;</span>
            <span class="c1"># NB: I copied some lines from self.compiler (which is an instance</span>
            <span class="c1"># of distutils.UnixCCompiler). See the following link.</span>
            <span class="c1"># https://github.com/python/cpython/blob/f03a8f8d5001963ad5b5b28dbd95497e9cc15596/Lib/distutils/ccompiler.py#L564-L567</span>
            <span class="c1"># This can be fragile, but a lot of other repos also do this</span>
            <span class="c1"># (see https://github.com/search?q=_setup_compile&amp;type=Code)</span>
            <span class="c1"># so it is probably OK; we&#39;ll also get CI signal if/when</span>
            <span class="c1"># we update our python version (which is when distutils can be</span>
            <span class="c1"># upgraded)</span>

            <span class="c1"># Use absolute path for output_dir so that the object file paths</span>
            <span class="c1"># (`objects`) get generated with absolute paths.</span>
            <span class="n">output_dir</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="n">output_dir</span><span class="p">)</span>
            <span class="n">_</span><span class="p">,</span> <span class="n">objects</span><span class="p">,</span> <span class="n">extra_postargs</span><span class="p">,</span> <span class="n">pp_opts</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> \
                <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">_setup_compile</span><span class="p">(</span><span class="n">output_dir</span><span class="p">,</span> <span class="n">macros</span><span class="p">,</span>
                                             <span class="n">include_dirs</span><span class="p">,</span> <span class="n">sources</span><span class="p">,</span>
                                             <span class="n">depends</span><span class="p">,</span> <span class="n">extra_postargs</span><span class="p">)</span>
            <span class="n">common_cflags</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">_get_cc_args</span><span class="p">(</span><span class="n">pp_opts</span><span class="p">,</span> <span class="n">debug</span><span class="p">,</span> <span class="n">extra_preargs</span><span class="p">)</span>
            <span class="n">extra_cc_cflags</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compiler_so</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span>
            <span class="n">with_cuda</span> <span class="o">=</span> <span class="nb">any</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="n">_is_cuda_file</span><span class="p">,</span> <span class="n">sources</span><span class="p">))</span>

            <span class="c1"># extra_postargs can be either:</span>
            <span class="c1"># - a dict mapping cxx/nvcc to extra flags</span>
            <span class="c1"># - a list of extra flags.</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">extra_postargs</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
                <span class="n">post_cflags</span> <span class="o">=</span> <span class="n">extra_postargs</span><span class="p">[</span><span class="s1">&#39;cxx&#39;</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">post_cflags</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">extra_postargs</span><span class="p">)</span>
            <span class="n">append_std14_if_no_std_present</span><span class="p">(</span><span class="n">post_cflags</span><span class="p">)</span>

            <span class="n">cuda_post_cflags</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">cuda_cflags</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
                <span class="n">cuda_cflags</span> <span class="o">=</span> <span class="n">common_cflags</span>
                <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">extra_postargs</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
                    <span class="n">cuda_post_cflags</span> <span class="o">=</span> <span class="n">extra_postargs</span><span class="p">[</span><span class="s1">&#39;nvcc&#39;</span><span class="p">]</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">cuda_post_cflags</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">extra_postargs</span><span class="p">)</span>
                <span class="n">cuda_post_cflags</span> <span class="o">=</span> <span class="n">unix_cuda_flags</span><span class="p">(</span><span class="n">cuda_post_cflags</span><span class="p">)</span>
                <span class="n">append_std14_if_no_std_present</span><span class="p">(</span><span class="n">cuda_post_cflags</span><span class="p">)</span>

            <span class="n">_write_ninja_file_and_compile_objects</span><span class="p">(</span>
                <span class="n">sources</span><span class="o">=</span><span class="n">sources</span><span class="p">,</span>
                <span class="n">objects</span><span class="o">=</span><span class="n">objects</span><span class="p">,</span>
                <span class="n">cflags</span><span class="o">=</span><span class="n">extra_cc_cflags</span> <span class="o">+</span> <span class="n">common_cflags</span><span class="p">,</span>
                <span class="n">post_cflags</span><span class="o">=</span><span class="n">post_cflags</span><span class="p">,</span>
                <span class="n">cuda_cflags</span><span class="o">=</span><span class="n">cuda_cflags</span><span class="p">,</span>
                <span class="n">cuda_post_cflags</span><span class="o">=</span><span class="n">cuda_post_cflags</span><span class="p">,</span>
                <span class="n">build_directory</span><span class="o">=</span><span class="n">output_dir</span><span class="p">,</span>
                <span class="n">verbose</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                <span class="n">with_cuda</span><span class="o">=</span><span class="n">with_cuda</span><span class="p">)</span>

            <span class="c1"># Return *all* object filenames, not just the ones we just built.</span>
            <span class="k">return</span> <span class="n">objects</span>

        <span class="k">def</span> <span class="nf">win_cuda_flags</span><span class="p">(</span><span class="n">cflags</span><span class="p">):</span>
            <span class="k">return</span> <span class="p">(</span><span class="n">COMMON_NVCC_FLAGS</span> <span class="o">+</span>
                    <span class="n">cflags</span> <span class="o">+</span> <span class="n">_get_cuda_arch_flags</span><span class="p">(</span><span class="n">cflags</span><span class="p">))</span>

        <span class="k">def</span> <span class="nf">win_wrap_single_compile</span><span class="p">(</span><span class="n">sources</span><span class="p">,</span>
                                    <span class="n">output_dir</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                    <span class="n">macros</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                    <span class="n">include_dirs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                    <span class="n">debug</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                                    <span class="n">extra_preargs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                    <span class="n">extra_postargs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                    <span class="n">depends</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>

            <span class="bp">self</span><span class="o">.</span><span class="n">cflags</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">extra_postargs</span><span class="p">)</span>
            <span class="n">extra_postargs</span> <span class="o">=</span> <span class="kc">None</span>

            <span class="k">def</span> <span class="nf">spawn</span><span class="p">(</span><span class="n">cmd</span><span class="p">):</span>
                <span class="c1"># Using regex to match src, obj and include files</span>
                <span class="n">src_regex</span> <span class="o">=</span> <span class="n">re</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="s1">&#39;/T(p|c)(.*)&#39;</span><span class="p">)</span>
                <span class="n">src_list</span> <span class="o">=</span> <span class="p">[</span>
                    <span class="n">m</span><span class="o">.</span><span class="n">group</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span> <span class="k">for</span> <span class="n">m</span> <span class="ow">in</span> <span class="p">(</span><span class="n">src_regex</span><span class="o">.</span><span class="n">match</span><span class="p">(</span><span class="n">elem</span><span class="p">)</span> <span class="k">for</span> <span class="n">elem</span> <span class="ow">in</span> <span class="n">cmd</span><span class="p">)</span>
                    <span class="k">if</span> <span class="n">m</span>
                <span class="p">]</span>

                <span class="n">obj_regex</span> <span class="o">=</span> <span class="n">re</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="s1">&#39;/Fo(.*)&#39;</span><span class="p">)</span>
                <span class="n">obj_list</span> <span class="o">=</span> <span class="p">[</span>
                    <span class="n">m</span><span class="o">.</span><span class="n">group</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span> <span class="k">for</span> <span class="n">m</span> <span class="ow">in</span> <span class="p">(</span><span class="n">obj_regex</span><span class="o">.</span><span class="n">match</span><span class="p">(</span><span class="n">elem</span><span class="p">)</span> <span class="k">for</span> <span class="n">elem</span> <span class="ow">in</span> <span class="n">cmd</span><span class="p">)</span>
                    <span class="k">if</span> <span class="n">m</span>
                <span class="p">]</span>

                <span class="n">include_regex</span> <span class="o">=</span> <span class="n">re</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="sa">r</span><span class="s1">&#39;((\-|\/)I.*)&#39;</span><span class="p">)</span>
                <span class="n">include_list</span> <span class="o">=</span> <span class="p">[</span>
                    <span class="n">m</span><span class="o">.</span><span class="n">group</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
                    <span class="k">for</span> <span class="n">m</span> <span class="ow">in</span> <span class="p">(</span><span class="n">include_regex</span><span class="o">.</span><span class="n">match</span><span class="p">(</span><span class="n">elem</span><span class="p">)</span> <span class="k">for</span> <span class="n">elem</span> <span class="ow">in</span> <span class="n">cmd</span><span class="p">)</span> <span class="k">if</span> <span class="n">m</span>
                <span class="p">]</span>

                <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">src_list</span><span class="p">)</span> <span class="o">&gt;=</span> <span class="mi">1</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">obj_list</span><span class="p">)</span> <span class="o">&gt;=</span> <span class="mi">1</span><span class="p">:</span>
                    <span class="n">src</span> <span class="o">=</span> <span class="n">src_list</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                    <span class="n">obj</span> <span class="o">=</span> <span class="n">obj_list</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                    <span class="k">if</span> <span class="n">_is_cuda_file</span><span class="p">(</span><span class="n">src</span><span class="p">):</span>
                        <span class="n">nvcc</span> <span class="o">=</span> <span class="n">_join_cuda_home</span><span class="p">(</span><span class="s1">&#39;bin&#39;</span><span class="p">,</span> <span class="s1">&#39;nvcc&#39;</span><span class="p">)</span>
                        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cflags</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
                            <span class="n">cflags</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cflags</span><span class="p">[</span><span class="s1">&#39;nvcc&#39;</span><span class="p">]</span>
                        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cflags</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
                            <span class="n">cflags</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cflags</span>
                        <span class="k">else</span><span class="p">:</span>
                            <span class="n">cflags</span> <span class="o">=</span> <span class="p">[]</span>

                        <span class="n">cflags</span> <span class="o">=</span> <span class="n">win_cuda_flags</span><span class="p">(</span><span class="n">cflags</span><span class="p">)</span>
                        <span class="k">for</span> <span class="n">flag</span> <span class="ow">in</span> <span class="n">COMMON_MSVC_FLAGS</span><span class="p">:</span>
                            <span class="n">cflags</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;-Xcompiler&#39;</span><span class="p">,</span> <span class="n">flag</span><span class="p">]</span> <span class="o">+</span> <span class="n">cflags</span>
                        <span class="n">cmd</span> <span class="o">=</span> <span class="p">[</span><span class="n">nvcc</span><span class="p">,</span> <span class="s1">&#39;-c&#39;</span><span class="p">,</span> <span class="n">src</span><span class="p">,</span> <span class="s1">&#39;-o&#39;</span><span class="p">,</span> <span class="n">obj</span><span class="p">]</span> <span class="o">+</span> <span class="n">include_list</span> <span class="o">+</span> <span class="n">cflags</span>
                    <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cflags</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
                        <span class="n">cflags</span> <span class="o">=</span> <span class="n">COMMON_MSVC_FLAGS</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">cflags</span><span class="p">[</span><span class="s1">&#39;cxx&#39;</span><span class="p">]</span>
                        <span class="n">cmd</span> <span class="o">+=</span> <span class="n">cflags</span>
                    <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cflags</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
                        <span class="n">cflags</span> <span class="o">=</span> <span class="n">COMMON_MSVC_FLAGS</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">cflags</span>
                        <span class="n">cmd</span> <span class="o">+=</span> <span class="n">cflags</span>

                <span class="k">return</span> <span class="n">original_spawn</span><span class="p">(</span><span class="n">cmd</span><span class="p">)</span>

            <span class="k">try</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">spawn</span> <span class="o">=</span> <span class="n">spawn</span>
                <span class="k">return</span> <span class="n">original_compile</span><span class="p">(</span><span class="n">sources</span><span class="p">,</span> <span class="n">output_dir</span><span class="p">,</span> <span class="n">macros</span><span class="p">,</span>
                                        <span class="n">include_dirs</span><span class="p">,</span> <span class="n">debug</span><span class="p">,</span> <span class="n">extra_preargs</span><span class="p">,</span>
                                        <span class="n">extra_postargs</span><span class="p">,</span> <span class="n">depends</span><span class="p">)</span>
            <span class="k">finally</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">spawn</span> <span class="o">=</span> <span class="n">original_spawn</span>

        <span class="k">def</span> <span class="nf">win_wrap_ninja_compile</span><span class="p">(</span><span class="n">sources</span><span class="p">,</span>
                                   <span class="n">output_dir</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                   <span class="n">macros</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                   <span class="n">include_dirs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                   <span class="n">debug</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                                   <span class="n">extra_preargs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                   <span class="n">extra_postargs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                                   <span class="n">depends</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>

            <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">initialized</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">initialize</span><span class="p">()</span>
            <span class="n">output_dir</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="n">output_dir</span><span class="p">)</span>
            <span class="n">_</span><span class="p">,</span> <span class="n">objects</span><span class="p">,</span> <span class="n">extra_postargs</span><span class="p">,</span> <span class="n">pp_opts</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> \
                <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">_setup_compile</span><span class="p">(</span><span class="n">output_dir</span><span class="p">,</span> <span class="n">macros</span><span class="p">,</span>
                                             <span class="n">include_dirs</span><span class="p">,</span> <span class="n">sources</span><span class="p">,</span>
                                             <span class="n">depends</span><span class="p">,</span> <span class="n">extra_postargs</span><span class="p">)</span>
            <span class="n">common_cflags</span> <span class="o">=</span> <span class="n">extra_preargs</span> <span class="ow">or</span> <span class="p">[]</span>
            <span class="n">cflags</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="k">if</span> <span class="n">debug</span><span class="p">:</span>
                <span class="n">cflags</span><span class="o">.</span><span class="n">extend</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compile_options_debug</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">cflags</span><span class="o">.</span><span class="n">extend</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compile_options</span><span class="p">)</span>
            <span class="n">common_cflags</span><span class="o">.</span><span class="n">extend</span><span class="p">(</span><span class="n">COMMON_MSVC_FLAGS</span><span class="p">)</span>
            <span class="n">cflags</span> <span class="o">=</span> <span class="n">cflags</span> <span class="o">+</span> <span class="n">common_cflags</span> <span class="o">+</span> <span class="n">pp_opts</span>
            <span class="n">with_cuda</span> <span class="o">=</span> <span class="nb">any</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="n">_is_cuda_file</span><span class="p">,</span> <span class="n">sources</span><span class="p">))</span>

            <span class="c1"># extra_postargs can be either:</span>
            <span class="c1"># - a dict mapping cxx/nvcc to extra flags</span>
            <span class="c1"># - a list of extra flags.</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">extra_postargs</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
                <span class="n">post_cflags</span> <span class="o">=</span> <span class="n">extra_postargs</span><span class="p">[</span><span class="s1">&#39;cxx&#39;</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">post_cflags</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">extra_postargs</span><span class="p">)</span>
            <span class="n">append_std14_if_no_std_present</span><span class="p">(</span><span class="n">post_cflags</span><span class="p">)</span>

            <span class="n">cuda_post_cflags</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">cuda_cflags</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
                <span class="n">cuda_cflags</span> <span class="o">=</span> <span class="p">[]</span>
                <span class="k">for</span> <span class="n">common_cflag</span> <span class="ow">in</span> <span class="n">common_cflags</span><span class="p">:</span>
                    <span class="n">cuda_cflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-Xcompiler&#39;</span><span class="p">)</span>
                    <span class="n">cuda_cflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">common_cflag</span><span class="p">)</span>
                <span class="n">cuda_cflags</span><span class="o">.</span><span class="n">extend</span><span class="p">(</span><span class="n">pp_opts</span><span class="p">)</span>
                <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">extra_postargs</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
                    <span class="n">cuda_post_cflags</span> <span class="o">=</span> <span class="n">extra_postargs</span><span class="p">[</span><span class="s1">&#39;nvcc&#39;</span><span class="p">]</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">cuda_post_cflags</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">extra_postargs</span><span class="p">)</span>
                <span class="n">cuda_post_cflags</span> <span class="o">=</span> <span class="n">win_cuda_flags</span><span class="p">(</span><span class="n">cuda_post_cflags</span><span class="p">)</span>

            <span class="kn">from</span> <span class="nn">distutils.spawn</span> <span class="kn">import</span> <span class="n">_nt_quote_args</span>
            <span class="n">cflags</span> <span class="o">=</span> <span class="n">_nt_quote_args</span><span class="p">(</span><span class="n">cflags</span><span class="p">)</span>
            <span class="n">post_cflags</span> <span class="o">=</span> <span class="n">_nt_quote_args</span><span class="p">(</span><span class="n">post_cflags</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
                <span class="n">cuda_cflags</span> <span class="o">=</span> <span class="n">_nt_quote_args</span><span class="p">(</span><span class="n">cuda_cflags</span><span class="p">)</span>
                <span class="n">cuda_post_cflags</span> <span class="o">=</span> <span class="n">_nt_quote_args</span><span class="p">(</span><span class="n">cuda_post_cflags</span><span class="p">)</span>

            <span class="n">_write_ninja_file_and_compile_objects</span><span class="p">(</span>
                <span class="n">sources</span><span class="o">=</span><span class="n">sources</span><span class="p">,</span>
                <span class="n">objects</span><span class="o">=</span><span class="n">objects</span><span class="p">,</span>
                <span class="n">cflags</span><span class="o">=</span><span class="n">cflags</span><span class="p">,</span>
                <span class="n">post_cflags</span><span class="o">=</span><span class="n">post_cflags</span><span class="p">,</span>
                <span class="n">cuda_cflags</span><span class="o">=</span><span class="n">cuda_cflags</span><span class="p">,</span>
                <span class="n">cuda_post_cflags</span><span class="o">=</span><span class="n">cuda_post_cflags</span><span class="p">,</span>
                <span class="n">build_directory</span><span class="o">=</span><span class="n">output_dir</span><span class="p">,</span>
                <span class="n">verbose</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                <span class="n">with_cuda</span><span class="o">=</span><span class="n">with_cuda</span><span class="p">)</span>

            <span class="c1"># Return *all* object filenames, not just the ones we just built.</span>
            <span class="k">return</span> <span class="n">objects</span>

        <span class="c1"># Monkey-patch the _compile or compile method.</span>
        <span class="c1"># https://github.com/python/cpython/blob/dc0284ee8f7a270b6005467f26d8e5773d76e959/Lib/distutils/ccompiler.py#L511</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compiler_type</span> <span class="o">==</span> <span class="s1">&#39;msvc&#39;</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">use_ninja</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compile</span> <span class="o">=</span> <span class="n">win_wrap_ninja_compile</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compile</span> <span class="o">=</span> <span class="n">win_wrap_single_compile</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">use_ninja</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compile</span> <span class="o">=</span> <span class="n">unix_wrap_ninja_compile</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">_compile</span> <span class="o">=</span> <span class="n">unix_wrap_single_compile</span>

        <span class="n">build_ext</span><span class="o">.</span><span class="n">build_extensions</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">get_ext_filename</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ext_name</span><span class="p">):</span>
        <span class="c1"># Get the original shared library name. For Python 3, this name will be</span>
        <span class="c1"># suffixed with &quot;&lt;SOABI&gt;.so&quot;, where &lt;SOABI&gt; will be something like</span>
        <span class="c1"># cpython-37m-x86_64-linux-gnu. On Python 2, there is no such ABI name.</span>
        <span class="c1"># The final extension, .so, would be .lib/.dll on Windows of course.</span>
        <span class="n">ext_filename</span> <span class="o">=</span> <span class="nb">super</span><span class="p">(</span><span class="n">BuildExtension</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">get_ext_filename</span><span class="p">(</span><span class="n">ext_name</span><span class="p">)</span>
        <span class="c1"># If `no_python_abi_suffix` is `True`, we omit the Python 3 ABI</span>
        <span class="c1"># component. This makes building shared libraries with setuptools that</span>
        <span class="c1"># aren&#39;t Python modules nicer.</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">no_python_abi_suffix</span> <span class="ow">and</span> <span class="n">sys</span><span class="o">.</span><span class="n">version_info</span> <span class="o">&gt;=</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">0</span><span class="p">):</span>
            <span class="c1"># The parts will be e.g. [&quot;my_extension&quot;, &quot;cpython-37m-x86_64-linux-gnu&quot;, &quot;so&quot;].</span>
            <span class="n">ext_filename_parts</span> <span class="o">=</span> <span class="n">ext_filename</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s1">&#39;.&#39;</span><span class="p">)</span>
            <span class="c1"># Omit the second to last element.</span>
            <span class="n">without_abi</span> <span class="o">=</span> <span class="n">ext_filename_parts</span><span class="p">[:</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="n">ext_filename_parts</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">:]</span>
            <span class="n">ext_filename</span> <span class="o">=</span> <span class="s1">&#39;.&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">without_abi</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">ext_filename</span>

    <span class="k">def</span> <span class="nf">_check_abi</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="c1"># On some platforms, like Windows, compiler_cxx is not available.</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="p">,</span> <span class="s1">&#39;compiler_cxx&#39;</span><span class="p">):</span>
            <span class="n">compiler</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler</span><span class="o">.</span><span class="n">compiler_cxx</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="k">elif</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
            <span class="n">compiler</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CXX&#39;</span><span class="p">,</span> <span class="s1">&#39;cl&#39;</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">compiler</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CXX&#39;</span><span class="p">,</span> <span class="s1">&#39;c++&#39;</span><span class="p">)</span>
        <span class="n">check_compiler_abi_compatibility</span><span class="p">(</span><span class="n">compiler</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_add_compile_flag</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">extension</span><span class="p">,</span> <span class="n">flag</span><span class="p">):</span>
        <span class="n">extension</span><span class="o">.</span><span class="n">extra_compile_args</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">extension</span><span class="o">.</span><span class="n">extra_compile_args</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">extension</span><span class="o">.</span><span class="n">extra_compile_args</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="k">for</span> <span class="n">args</span> <span class="ow">in</span> <span class="n">extension</span><span class="o">.</span><span class="n">extra_compile_args</span><span class="o">.</span><span class="n">values</span><span class="p">():</span>
                <span class="n">args</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">flag</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">extension</span><span class="o">.</span><span class="n">extra_compile_args</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">flag</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_define_torch_extension_name</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">extension</span><span class="p">):</span>
        <span class="c1"># pybind11 doesn&#39;t support dots in the names</span>
        <span class="c1"># so in order to support extensions in the packages</span>
        <span class="c1"># like torch._C, we take the last part of the string</span>
        <span class="c1"># as the library name</span>
        <span class="n">names</span> <span class="o">=</span> <span class="n">extension</span><span class="o">.</span><span class="n">name</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s1">&#39;.&#39;</span><span class="p">)</span>
        <span class="n">name</span> <span class="o">=</span> <span class="n">names</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
        <span class="n">define</span> <span class="o">=</span> <span class="s1">&#39;-DTORCH_EXTENSION_NAME=</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_add_compile_flag</span><span class="p">(</span><span class="n">extension</span><span class="p">,</span> <span class="n">define</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_add_gnu_cpp_abi_flag</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">extension</span><span class="p">):</span>
        <span class="c1"># use the same CXX ABI as what PyTorch was compiled with</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_add_compile_flag</span><span class="p">(</span><span class="n">extension</span><span class="p">,</span> <span class="s1">&#39;-D_GLIBCXX_USE_CXX11_ABI=&#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_GLIBCXX_USE_CXX11_ABI</span><span class="p">)))</span></div>


<div class="viewcode-block" id="CppExtension"><a class="viewcode-back" href="../../../cpp_extension.html#torch.utils.cpp_extension.CppExtension">[docs]</a><span class="k">def</span> <span class="nf">CppExtension</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">sources</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Creates a :class:`setuptools.Extension` for C++.</span>

<span class="sd">    Convenience method that creates a :class:`setuptools.Extension` with the</span>
<span class="sd">    bare minimum (but often sufficient) arguments to build a C++ extension.</span>

<span class="sd">    All arguments are forwarded to the :class:`setuptools.Extension`</span>
<span class="sd">    constructor.</span>

<span class="sd">    Example:</span>
<span class="sd">        &gt;&gt;&gt; from setuptools import setup</span>
<span class="sd">        &gt;&gt;&gt; from torch.utils.cpp_extension import BuildExtension, CppExtension</span>
<span class="sd">        &gt;&gt;&gt; setup(</span>
<span class="sd">                name=&#39;extension&#39;,</span>
<span class="sd">                ext_modules=[</span>
<span class="sd">                    CppExtension(</span>
<span class="sd">                        name=&#39;extension&#39;,</span>
<span class="sd">                        sources=[&#39;extension.cpp&#39;],</span>
<span class="sd">                        extra_compile_args=[&#39;-g&#39;]),</span>
<span class="sd">                ],</span>
<span class="sd">                cmdclass={</span>
<span class="sd">                    &#39;build_ext&#39;: BuildExtension</span>
<span class="sd">                })</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="n">include_dirs</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;include_dirs&#39;</span><span class="p">,</span> <span class="p">[])</span>
    <span class="n">include_dirs</span> <span class="o">+=</span> <span class="n">include_paths</span><span class="p">()</span>
    <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;include_dirs&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">include_dirs</span>

    <span class="n">library_dirs</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;library_dirs&#39;</span><span class="p">,</span> <span class="p">[])</span>
    <span class="n">library_dirs</span> <span class="o">+=</span> <span class="n">library_paths</span><span class="p">()</span>
    <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;library_dirs&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">library_dirs</span>

    <span class="n">libraries</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;libraries&#39;</span><span class="p">,</span> <span class="p">[])</span>
    <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;c10&#39;</span><span class="p">)</span>
    <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch&#39;</span><span class="p">)</span>
    <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch_cpu&#39;</span><span class="p">)</span>
    <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch_python&#39;</span><span class="p">)</span>
    <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;libraries&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">libraries</span>

    <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;language&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="s1">&#39;c++&#39;</span>
    <span class="k">return</span> <span class="n">setuptools</span><span class="o">.</span><span class="n">Extension</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">sources</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>


<div class="viewcode-block" id="CUDAExtension"><a class="viewcode-back" href="../../../cpp_extension.html#torch.utils.cpp_extension.CUDAExtension">[docs]</a><span class="k">def</span> <span class="nf">CUDAExtension</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">sources</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Creates a :class:`setuptools.Extension` for CUDA/C++.</span>

<span class="sd">    Convenience method that creates a :class:`setuptools.Extension` with the</span>
<span class="sd">    bare minimum (but often sufficient) arguments to build a CUDA/C++</span>
<span class="sd">    extension. This includes the CUDA include path, library path and runtime</span>
<span class="sd">    library.</span>

<span class="sd">    All arguments are forwarded to the :class:`setuptools.Extension`</span>
<span class="sd">    constructor.</span>

<span class="sd">    Example:</span>
<span class="sd">        &gt;&gt;&gt; from setuptools import setup</span>
<span class="sd">        &gt;&gt;&gt; from torch.utils.cpp_extension import BuildExtension, CUDAExtension</span>
<span class="sd">        &gt;&gt;&gt; setup(</span>
<span class="sd">                name=&#39;cuda_extension&#39;,</span>
<span class="sd">                ext_modules=[</span>
<span class="sd">                    CUDAExtension(</span>
<span class="sd">                            name=&#39;cuda_extension&#39;,</span>
<span class="sd">                            sources=[&#39;extension.cpp&#39;, &#39;extension_kernel.cu&#39;],</span>
<span class="sd">                            extra_compile_args={&#39;cxx&#39;: [&#39;-g&#39;],</span>
<span class="sd">                                                &#39;nvcc&#39;: [&#39;-O2&#39;]})</span>
<span class="sd">                ],</span>
<span class="sd">                cmdclass={</span>
<span class="sd">                    &#39;build_ext&#39;: BuildExtension</span>
<span class="sd">                })</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="n">library_dirs</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;library_dirs&#39;</span><span class="p">,</span> <span class="p">[])</span>
    <span class="n">library_dirs</span> <span class="o">+=</span> <span class="n">library_paths</span><span class="p">(</span><span class="n">cuda</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
    <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;library_dirs&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">library_dirs</span>

    <span class="n">libraries</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;libraries&#39;</span><span class="p">,</span> <span class="p">[])</span>
    <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;c10&#39;</span><span class="p">)</span>
    <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch&#39;</span><span class="p">)</span>
    <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch_cpu&#39;</span><span class="p">)</span>
    <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch_python&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">IS_HIP_EXTENSION</span><span class="p">:</span>
        <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;c10_hip&#39;</span><span class="p">)</span>
        <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch_hip&#39;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;cudart&#39;</span><span class="p">)</span>
        <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;c10_cuda&#39;</span><span class="p">)</span>
        <span class="n">libraries</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch_cuda&#39;</span><span class="p">)</span>
    <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;libraries&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">libraries</span>

    <span class="n">include_dirs</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;include_dirs&#39;</span><span class="p">,</span> <span class="p">[])</span>
    <span class="n">include_dirs</span> <span class="o">+=</span> <span class="n">include_paths</span><span class="p">(</span><span class="n">cuda</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
    <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;include_dirs&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">include_dirs</span>

    <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;language&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="s1">&#39;c++&#39;</span>

    <span class="k">return</span> <span class="n">setuptools</span><span class="o">.</span><span class="n">Extension</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">sources</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>


<div class="viewcode-block" id="include_paths"><a class="viewcode-back" href="../../../cpp_extension.html#torch.utils.cpp_extension.include_paths">[docs]</a><span class="k">def</span> <span class="nf">include_paths</span><span class="p">(</span><span class="n">cuda</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Get the include paths required to build a C++ or CUDA extension.</span>

<span class="sd">    Args:</span>
<span class="sd">        cuda: If `True`, includes CUDA-specific include paths.</span>

<span class="sd">    Returns:</span>
<span class="sd">        A list of include path strings.</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="n">here</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="vm">__file__</span><span class="p">)</span>
    <span class="n">torch_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">here</span><span class="p">))</span>
    <span class="n">lib_include</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">torch_path</span><span class="p">,</span> <span class="s1">&#39;include&#39;</span><span class="p">)</span>
    <span class="n">paths</span> <span class="o">=</span> <span class="p">[</span>
        <span class="n">lib_include</span><span class="p">,</span>
        <span class="c1"># Remove this once torch/torch.h is officially no longer supported for C++ extensions.</span>
        <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">lib_include</span><span class="p">,</span> <span class="s1">&#39;torch&#39;</span><span class="p">,</span> <span class="s1">&#39;csrc&#39;</span><span class="p">,</span> <span class="s1">&#39;api&#39;</span><span class="p">,</span> <span class="s1">&#39;include&#39;</span><span class="p">),</span>
        <span class="c1"># Some internal (old) Torch headers don&#39;t properly prefix their includes,</span>
        <span class="c1"># so we need to pass -Itorch/lib/include/TH as well.</span>
        <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">lib_include</span><span class="p">,</span> <span class="s1">&#39;TH&#39;</span><span class="p">),</span>
        <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">lib_include</span><span class="p">,</span> <span class="s1">&#39;THC&#39;</span><span class="p">)</span>
    <span class="p">]</span>
    <span class="k">if</span> <span class="n">cuda</span> <span class="ow">and</span> <span class="n">IS_HIP_EXTENSION</span><span class="p">:</span>
        <span class="n">paths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">lib_include</span><span class="p">,</span> <span class="s1">&#39;THH&#39;</span><span class="p">))</span>
        <span class="n">paths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">_join_rocm_home</span><span class="p">(</span><span class="s1">&#39;include&#39;</span><span class="p">))</span>
        <span class="k">if</span> <span class="n">MIOPEN_HOME</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">paths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">MIOPEN_HOME</span><span class="p">,</span> <span class="s1">&#39;include&#39;</span><span class="p">))</span>
    <span class="k">elif</span> <span class="n">cuda</span><span class="p">:</span>
        <span class="n">cuda_home_include</span> <span class="o">=</span> <span class="n">_join_cuda_home</span><span class="p">(</span><span class="s1">&#39;include&#39;</span><span class="p">)</span>
        <span class="c1"># if we have the Debian/Ubuntu packages for cuda, we get /usr as cuda home.</span>
        <span class="c1"># but gcc doesn&#39;t like having /usr/include passed explicitly</span>
        <span class="k">if</span> <span class="n">cuda_home_include</span> <span class="o">!=</span> <span class="s1">&#39;/usr/include&#39;</span><span class="p">:</span>
            <span class="n">paths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">cuda_home_include</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">CUDNN_HOME</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">paths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">CUDNN_HOME</span><span class="p">,</span> <span class="s1">&#39;include&#39;</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">paths</span></div>


<span class="k">def</span> <span class="nf">library_paths</span><span class="p">(</span><span class="n">cuda</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Get the library paths required to build a C++ or CUDA extension.</span>

<span class="sd">    Args:</span>
<span class="sd">        cuda: If `True`, includes CUDA-specific library paths.</span>

<span class="sd">    Returns:</span>
<span class="sd">        A list of library path strings.</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="n">paths</span> <span class="o">=</span> <span class="p">[]</span>

    <span class="c1"># We need to link against libtorch.so</span>
    <span class="n">here</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="vm">__file__</span><span class="p">)</span>
    <span class="n">torch_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">here</span><span class="p">))</span>
    <span class="n">lib_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">torch_path</span><span class="p">,</span> <span class="s1">&#39;lib&#39;</span><span class="p">)</span>
    <span class="n">paths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">lib_path</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">cuda</span> <span class="ow">and</span> <span class="n">IS_HIP_EXTENSION</span><span class="p">:</span>
        <span class="n">lib_dir</span> <span class="o">=</span> <span class="s1">&#39;lib&#39;</span>
        <span class="n">paths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">_join_rocm_home</span><span class="p">(</span><span class="n">lib_dir</span><span class="p">))</span>
    <span class="k">elif</span> <span class="n">cuda</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
            <span class="n">lib_dir</span> <span class="o">=</span> <span class="s1">&#39;lib/x64&#39;</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">lib_dir</span> <span class="o">=</span> <span class="s1">&#39;lib64&#39;</span>
            <span class="k">if</span> <span class="p">(</span><span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">_join_cuda_home</span><span class="p">(</span><span class="n">lib_dir</span><span class="p">))</span> <span class="ow">and</span>
                    <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">_join_cuda_home</span><span class="p">(</span><span class="s1">&#39;lib&#39;</span><span class="p">))):</span>
                <span class="c1"># 64-bit CUDA may be installed in &#39;lib&#39; (see e.g. gh-16955)</span>
                <span class="c1"># Note that it&#39;s also possible both don&#39;t exist (see</span>
                <span class="c1"># _find_cuda_home) - in that case we stay with &#39;lib64&#39;.</span>
                <span class="n">lib_dir</span> <span class="o">=</span> <span class="s1">&#39;lib&#39;</span>

        <span class="n">paths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">_join_cuda_home</span><span class="p">(</span><span class="n">lib_dir</span><span class="p">))</span>
        <span class="k">if</span> <span class="n">CUDNN_HOME</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">paths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">CUDNN_HOME</span><span class="p">,</span> <span class="n">lib_dir</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">paths</span>


<div class="viewcode-block" id="load"><a class="viewcode-back" href="../../../cpp_extension.html#torch.utils.cpp_extension.load">[docs]</a><span class="k">def</span> <span class="nf">load</span><span class="p">(</span><span class="n">name</span><span class="p">,</span>
         <span class="n">sources</span><span class="p">,</span>
         <span class="n">extra_cflags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
         <span class="n">extra_cuda_cflags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
         <span class="n">extra_ldflags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
         <span class="n">extra_include_paths</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
         <span class="n">build_directory</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
         <span class="n">verbose</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
         <span class="n">with_cuda</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
         <span class="n">is_python_module</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Loads a PyTorch C++ extension just-in-time (JIT).</span>

<span class="sd">    To load an extension, a Ninja build file is emitted, which is used to</span>
<span class="sd">    compile the given sources into a dynamic library. This library is</span>
<span class="sd">    subsequently loaded into the current Python process as a module and</span>
<span class="sd">    returned from this function, ready for use.</span>

<span class="sd">    By default, the directory to which the build file is emitted and the</span>
<span class="sd">    resulting library compiled to is ``&lt;tmp&gt;/torch_extensions/&lt;name&gt;``, where</span>
<span class="sd">    ``&lt;tmp&gt;`` is the temporary folder on the current platform and ``&lt;name&gt;``</span>
<span class="sd">    the name of the extension. This location can be overridden in two ways.</span>
<span class="sd">    First, if the ``TORCH_EXTENSIONS_DIR`` environment variable is set, it</span>
<span class="sd">    replaces ``&lt;tmp&gt;/torch_extensions`` and all extensions will be compiled</span>
<span class="sd">    into subfolders of this directory. Second, if the ``build_directory``</span>
<span class="sd">    argument to this function is supplied, it overrides the entire path, i.e.</span>
<span class="sd">    the library will be compiled into that folder directly.</span>

<span class="sd">    To compile the sources, the default system compiler (``c++``) is used,</span>
<span class="sd">    which can be overridden by setting the ``CXX`` environment variable. To pass</span>
<span class="sd">    additional arguments to the compilation process, ``extra_cflags`` or</span>
<span class="sd">    ``extra_ldflags`` can be provided. For example, to compile your extension</span>
<span class="sd">    with optimizations, pass ``extra_cflags=[&#39;-O3&#39;]``. You can also use</span>
<span class="sd">    ``extra_cflags`` to pass further include directories.</span>

<span class="sd">    CUDA support with mixed compilation is provided. Simply pass CUDA source</span>
<span class="sd">    files (``.cu`` or ``.cuh``) along with other sources. Such files will be</span>
<span class="sd">    detected and compiled with nvcc rather than the C++ compiler. This includes</span>
<span class="sd">    passing the CUDA lib64 directory as a library directory, and linking</span>
<span class="sd">    ``cudart``. You can pass additional flags to nvcc via</span>
<span class="sd">    ``extra_cuda_cflags``, just like with ``extra_cflags`` for C++. Various</span>
<span class="sd">    heuristics for finding the CUDA install directory are used, which usually</span>
<span class="sd">    work fine. If not, setting the ``CUDA_HOME`` environment variable is the</span>
<span class="sd">    safest option.</span>

<span class="sd">    Args:</span>
<span class="sd">        name: The name of the extension to build. This MUST be the same as the</span>
<span class="sd">            name of the pybind11 module!</span>
<span class="sd">        sources: A list of relative or absolute paths to C++ source files.</span>
<span class="sd">        extra_cflags: optional list of compiler flags to forward to the build.</span>
<span class="sd">        extra_cuda_cflags: optional list of compiler flags to forward to nvcc</span>
<span class="sd">            when building CUDA sources.</span>
<span class="sd">        extra_ldflags: optional list of linker flags to forward to the build.</span>
<span class="sd">        extra_include_paths: optional list of include directories to forward</span>
<span class="sd">            to the build.</span>
<span class="sd">        build_directory: optional path to use as build workspace.</span>
<span class="sd">        verbose: If ``True``, turns on verbose logging of load steps.</span>
<span class="sd">        with_cuda: Determines whether CUDA headers and libraries are added to</span>
<span class="sd">            the build. If set to ``None`` (default), this value is</span>
<span class="sd">            automatically determined based on the existence of ``.cu`` or</span>
<span class="sd">            ``.cuh`` in ``sources``. Set it to `True`` to force CUDA headers</span>
<span class="sd">            and libraries to be included.</span>
<span class="sd">        is_python_module: If ``True`` (default), imports the produced shared</span>
<span class="sd">            library as a Python module. If ``False``, loads it into the process</span>
<span class="sd">            as a plain dynamic library.</span>

<span class="sd">    Returns:</span>
<span class="sd">        If ``is_python_module`` is ``True``, returns the loaded PyTorch</span>
<span class="sd">        extension as a Python module. If ``is_python_module`` is ``False``</span>
<span class="sd">        returns nothing (the shared library is loaded into the process as a side</span>
<span class="sd">        effect).</span>

<span class="sd">    Example:</span>
<span class="sd">        &gt;&gt;&gt; from torch.utils.cpp_extension import load</span>
<span class="sd">        &gt;&gt;&gt; module = load(</span>
<span class="sd">                name=&#39;extension&#39;,</span>
<span class="sd">                sources=[&#39;extension.cpp&#39;, &#39;extension_kernel.cu&#39;],</span>
<span class="sd">                extra_cflags=[&#39;-O2&#39;],</span>
<span class="sd">                verbose=True)</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="k">return</span> <span class="n">_jit_compile</span><span class="p">(</span>
        <span class="n">name</span><span class="p">,</span>
        <span class="p">[</span><span class="n">sources</span><span class="p">]</span> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">sources</span><span class="p">,</span> <span class="nb">str</span><span class="p">)</span> <span class="k">else</span> <span class="n">sources</span><span class="p">,</span>
        <span class="n">extra_cflags</span><span class="p">,</span>
        <span class="n">extra_cuda_cflags</span><span class="p">,</span>
        <span class="n">extra_ldflags</span><span class="p">,</span>
        <span class="n">extra_include_paths</span><span class="p">,</span>
        <span class="n">build_directory</span> <span class="ow">or</span> <span class="n">_get_build_directory</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">verbose</span><span class="p">),</span>
        <span class="n">verbose</span><span class="p">,</span>
        <span class="n">with_cuda</span><span class="p">,</span>
        <span class="n">is_python_module</span><span class="p">)</span></div>


<div class="viewcode-block" id="load_inline"><a class="viewcode-back" href="../../../cpp_extension.html#torch.utils.cpp_extension.load_inline">[docs]</a><span class="k">def</span> <span class="nf">load_inline</span><span class="p">(</span><span class="n">name</span><span class="p">,</span>
                <span class="n">cpp_sources</span><span class="p">,</span>
                <span class="n">cuda_sources</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">functions</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">extra_cflags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">extra_cuda_cflags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">extra_ldflags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">extra_include_paths</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">build_directory</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">verbose</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                <span class="n">with_cuda</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">is_python_module</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                <span class="n">with_pytorch_error_handling</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Loads a PyTorch C++ extension just-in-time (JIT) from string sources.</span>

<span class="sd">    This function behaves exactly like :func:`load`, but takes its sources as</span>
<span class="sd">    strings rather than filenames. These strings are stored to files in the</span>
<span class="sd">    build directory, after which the behavior of :func:`load_inline` is</span>
<span class="sd">    identical to :func:`load`.</span>

<span class="sd">    See `the</span>
<span class="sd">    tests &lt;https://github.com/pytorch/pytorch/blob/master/test/test_cpp_extensions.py&gt;`_</span>
<span class="sd">    for good examples of using this function.</span>

<span class="sd">    Sources may omit two required parts of a typical non-inline C++ extension:</span>
<span class="sd">    the necessary header includes, as well as the (pybind11) binding code. More</span>
<span class="sd">    precisely, strings passed to ``cpp_sources`` are first concatenated into a</span>
<span class="sd">    single ``.cpp`` file. This file is then prepended with ``#include</span>
<span class="sd">    &lt;torch/extension.h&gt;``.</span>

<span class="sd">    Furthermore, if the ``functions`` argument is supplied, bindings will be</span>
<span class="sd">    automatically generated for each function specified. ``functions`` can</span>
<span class="sd">    either be a list of function names, or a dictionary mapping from function</span>
<span class="sd">    names to docstrings. If a list is given, the name of each function is used</span>
<span class="sd">    as its docstring.</span>

<span class="sd">    The sources in ``cuda_sources`` are concatenated into a separate ``.cu``</span>
<span class="sd">    file and  prepended with ``torch/types.h``, ``cuda.h`` and</span>
<span class="sd">    ``cuda_runtime.h`` includes. The ``.cpp`` and ``.cu`` files are compiled</span>
<span class="sd">    separately, but ultimately linked into a single library. Note that no</span>
<span class="sd">    bindings are generated for functions in ``cuda_sources`` per  se. To bind</span>
<span class="sd">    to a CUDA kernel, you must create a C++ function that calls it, and either</span>
<span class="sd">    declare or define this C++ function in one of the ``cpp_sources`` (and</span>
<span class="sd">    include its name in ``functions``).</span>

<span class="sd">    See :func:`load` for a description of arguments omitted below.</span>

<span class="sd">    Args:</span>
<span class="sd">        cpp_sources: A string, or list of strings, containing C++ source code.</span>
<span class="sd">        cuda_sources: A string, or list of strings, containing CUDA source code.</span>
<span class="sd">        functions: A list of function names for which to generate function</span>
<span class="sd">            bindings. If a dictionary is given, it should map function names to</span>
<span class="sd">            docstrings (which are otherwise just the function names).</span>
<span class="sd">        with_cuda: Determines whether CUDA headers and libraries are added to</span>
<span class="sd">            the build. If set to ``None`` (default), this value is</span>
<span class="sd">            automatically determined based on whether ``cuda_sources`` is</span>
<span class="sd">            provided. Set it to ``True`` to force CUDA headers</span>
<span class="sd">            and libraries to be included.</span>
<span class="sd">        with_pytorch_error_handling: Determines whether pytorch error and</span>
<span class="sd">            warning macros are handled by pytorch instead of pybind. To do</span>
<span class="sd">            this, each function ``foo`` is called via an intermediary ``_safe_foo``</span>
<span class="sd">            function. This redirection might cause issues in obscure cases</span>
<span class="sd">            of cpp. This flag should be set to ``False`` when this redirect</span>
<span class="sd">            causes issues.</span>

<span class="sd">    Example:</span>
<span class="sd">        &gt;&gt;&gt; from torch.utils.cpp_extension import load_inline</span>
<span class="sd">        &gt;&gt;&gt; source = \&#39;\&#39;\&#39;</span>
<span class="sd">        at::Tensor sin_add(at::Tensor x, at::Tensor y) {</span>
<span class="sd">          return x.sin() + y.sin();</span>
<span class="sd">        }</span>
<span class="sd">        \&#39;\&#39;\&#39;</span>
<span class="sd">        &gt;&gt;&gt; module = load_inline(name=&#39;inline_extension&#39;,</span>
<span class="sd">                                 cpp_sources=[source],</span>
<span class="sd">                                 functions=[&#39;sin_add&#39;])</span>

<span class="sd">    .. note::</span>
<span class="sd">        By default, the Ninja backend uses #CPUS + 2 workers to build the</span>
<span class="sd">        extension. This may use up too many resources on some systems. One</span>
<span class="sd">        can control the number of workers by setting the `MAX_JOBS` environment</span>
<span class="sd">        variable to a non-negative number.</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="n">build_directory</span> <span class="o">=</span> <span class="n">build_directory</span> <span class="ow">or</span> <span class="n">_get_build_directory</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">verbose</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">cpp_sources</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
        <span class="n">cpp_sources</span> <span class="o">=</span> <span class="p">[</span><span class="n">cpp_sources</span><span class="p">]</span>
    <span class="n">cuda_sources</span> <span class="o">=</span> <span class="n">cuda_sources</span> <span class="ow">or</span> <span class="p">[]</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">cuda_sources</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
        <span class="n">cuda_sources</span> <span class="o">=</span> <span class="p">[</span><span class="n">cuda_sources</span><span class="p">]</span>

    <span class="n">cpp_sources</span><span class="o">.</span><span class="n">insert</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="s1">&#39;#include &lt;torch/extension.h&gt;&#39;</span><span class="p">)</span>

    <span class="c1"># If `functions` is supplied, we create the pybind11 bindings for the user.</span>
    <span class="c1"># Here, `functions` is (or becomes, after some processing) a map from</span>
    <span class="c1"># function names to function docstrings.</span>
    <span class="k">if</span> <span class="n">functions</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">module_def</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="n">module_def</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">functions</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="n">functions</span> <span class="o">=</span> <span class="p">[</span><span class="n">functions</span><span class="p">]</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">functions</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="c1"># Make the function docstring the same as the function name.</span>
            <span class="n">functions</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">((</span><span class="n">f</span><span class="p">,</span> <span class="n">f</span><span class="p">)</span> <span class="k">for</span> <span class="n">f</span> <span class="ow">in</span> <span class="n">functions</span><span class="p">)</span>
        <span class="k">elif</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">functions</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                <span class="s2">&quot;Expected &#39;functions&#39; to be a list or dict, but was </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
                    <span class="nb">type</span><span class="p">(</span><span class="n">functions</span><span class="p">)))</span>
        <span class="k">for</span> <span class="n">function_name</span><span class="p">,</span> <span class="n">docstring</span> <span class="ow">in</span> <span class="n">functions</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
            <span class="k">if</span> <span class="n">with_pytorch_error_handling</span><span class="p">:</span>
                <span class="n">module_def</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
                    <span class="s1">&#39;m.def(&quot;</span><span class="si">{0}</span><span class="s1">&quot;, torch::wrap_pybind_function(</span><span class="si">{0}</span><span class="s1">), &quot;</span><span class="si">{1}</span><span class="s1">&quot;);&#39;</span>
                    <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">function_name</span><span class="p">,</span> <span class="n">docstring</span><span class="p">))</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">module_def</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;m.def(&quot;</span><span class="si">{0}</span><span class="s1">&quot;, </span><span class="si">{0}</span><span class="s1">, &quot;</span><span class="si">{1}</span><span class="s1">&quot;);&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">function_name</span><span class="p">,</span> <span class="n">docstring</span><span class="p">))</span>
        <span class="n">module_def</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;}&#39;</span><span class="p">)</span>
        <span class="n">cpp_sources</span> <span class="o">+=</span> <span class="n">module_def</span>

    <span class="n">cpp_source_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">build_directory</span><span class="p">,</span> <span class="s1">&#39;main.cpp&#39;</span><span class="p">)</span>
    <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">cpp_source_path</span><span class="p">,</span> <span class="s1">&#39;w&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">cpp_source_file</span><span class="p">:</span>
        <span class="n">cpp_source_file</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s1">&#39;</span><span class="se">\n</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">cpp_sources</span><span class="p">))</span>

    <span class="n">sources</span> <span class="o">=</span> <span class="p">[</span><span class="n">cpp_source_path</span><span class="p">]</span>

    <span class="k">if</span> <span class="n">cuda_sources</span><span class="p">:</span>
        <span class="n">cuda_sources</span><span class="o">.</span><span class="n">insert</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="s1">&#39;#include &lt;torch/types.h&gt;&#39;</span><span class="p">)</span>
        <span class="n">cuda_sources</span><span class="o">.</span><span class="n">insert</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="s1">&#39;#include &lt;cuda.h&gt;&#39;</span><span class="p">)</span>
        <span class="n">cuda_sources</span><span class="o">.</span><span class="n">insert</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="s1">&#39;#include &lt;cuda_runtime.h&gt;&#39;</span><span class="p">)</span>

        <span class="n">cuda_source_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">build_directory</span><span class="p">,</span> <span class="s1">&#39;cuda.cu&#39;</span><span class="p">)</span>
        <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">cuda_source_path</span><span class="p">,</span> <span class="s1">&#39;w&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">cuda_source_file</span><span class="p">:</span>
            <span class="n">cuda_source_file</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s1">&#39;</span><span class="se">\n</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">cuda_sources</span><span class="p">))</span>

        <span class="n">sources</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">cuda_source_path</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">_jit_compile</span><span class="p">(</span>
        <span class="n">name</span><span class="p">,</span>
        <span class="n">sources</span><span class="p">,</span>
        <span class="n">extra_cflags</span><span class="p">,</span>
        <span class="n">extra_cuda_cflags</span><span class="p">,</span>
        <span class="n">extra_ldflags</span><span class="p">,</span>
        <span class="n">extra_include_paths</span><span class="p">,</span>
        <span class="n">build_directory</span><span class="p">,</span>
        <span class="n">verbose</span><span class="p">,</span>
        <span class="n">with_cuda</span><span class="p">,</span>
        <span class="n">is_python_module</span><span class="p">)</span></div>


<span class="k">def</span> <span class="nf">_jit_compile</span><span class="p">(</span><span class="n">name</span><span class="p">,</span>
                 <span class="n">sources</span><span class="p">,</span>
                 <span class="n">extra_cflags</span><span class="p">,</span>
                 <span class="n">extra_cuda_cflags</span><span class="p">,</span>
                 <span class="n">extra_ldflags</span><span class="p">,</span>
                 <span class="n">extra_include_paths</span><span class="p">,</span>
                 <span class="n">build_directory</span><span class="p">,</span>
                 <span class="n">verbose</span><span class="p">,</span>
                 <span class="n">with_cuda</span><span class="p">,</span>
                 <span class="n">is_python_module</span><span class="p">):</span>
    <span class="n">old_version</span> <span class="o">=</span> <span class="n">JIT_EXTENSION_VERSIONER</span><span class="o">.</span><span class="n">get_version</span><span class="p">(</span><span class="n">name</span><span class="p">)</span>
    <span class="n">version</span> <span class="o">=</span> <span class="n">JIT_EXTENSION_VERSIONER</span><span class="o">.</span><span class="n">bump_version_if_changed</span><span class="p">(</span>
        <span class="n">name</span><span class="p">,</span>
        <span class="n">sources</span><span class="p">,</span>
        <span class="n">build_arguments</span><span class="o">=</span><span class="p">[</span><span class="n">extra_cflags</span><span class="p">,</span> <span class="n">extra_cuda_cflags</span><span class="p">,</span> <span class="n">extra_ldflags</span><span class="p">,</span> <span class="n">extra_include_paths</span><span class="p">],</span>
        <span class="n">build_directory</span><span class="o">=</span><span class="n">build_directory</span><span class="p">,</span>
        <span class="n">with_cuda</span><span class="o">=</span><span class="n">with_cuda</span>
    <span class="p">)</span>
    <span class="k">if</span> <span class="n">version</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">version</span> <span class="o">!=</span> <span class="n">old_version</span> <span class="ow">and</span> <span class="n">verbose</span><span class="p">:</span>
            <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;The input conditions for extension module </span><span class="si">{}</span><span class="s1"> have changed. &#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">)</span> <span class="o">+</span>
                  <span class="s1">&#39;Bumping to version </span><span class="si">{0}</span><span class="s1"> and re-building as </span><span class="si">{1}</span><span class="s1">_v</span><span class="si">{0}</span><span class="s1">...&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">version</span><span class="p">,</span> <span class="n">name</span><span class="p">))</span>
        <span class="n">name</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="si">{}</span><span class="s1">_v</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">version</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">version</span> <span class="o">!=</span> <span class="n">old_version</span><span class="p">:</span>
        <span class="n">baton</span> <span class="o">=</span> <span class="n">FileBaton</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">build_directory</span><span class="p">,</span> <span class="s1">&#39;lock&#39;</span><span class="p">))</span>
        <span class="k">if</span> <span class="n">baton</span><span class="o">.</span><span class="n">try_acquire</span><span class="p">():</span>
            <span class="k">try</span><span class="p">:</span>
                <span class="n">_write_ninja_file_and_build_library</span><span class="p">(</span>
                    <span class="n">name</span><span class="o">=</span><span class="n">name</span><span class="p">,</span>
                    <span class="n">sources</span><span class="o">=</span><span class="n">sources</span><span class="p">,</span>
                    <span class="n">extra_cflags</span><span class="o">=</span><span class="n">extra_cflags</span> <span class="ow">or</span> <span class="p">[],</span>
                    <span class="n">extra_cuda_cflags</span><span class="o">=</span><span class="n">extra_cuda_cflags</span> <span class="ow">or</span> <span class="p">[],</span>
                    <span class="n">extra_ldflags</span><span class="o">=</span><span class="n">extra_ldflags</span> <span class="ow">or</span> <span class="p">[],</span>
                    <span class="n">extra_include_paths</span><span class="o">=</span><span class="n">extra_include_paths</span> <span class="ow">or</span> <span class="p">[],</span>
                    <span class="n">build_directory</span><span class="o">=</span><span class="n">build_directory</span><span class="p">,</span>
                    <span class="n">verbose</span><span class="o">=</span><span class="n">verbose</span><span class="p">,</span>
                    <span class="n">with_cuda</span><span class="o">=</span><span class="n">with_cuda</span><span class="p">)</span>
            <span class="k">finally</span><span class="p">:</span>
                <span class="n">baton</span><span class="o">.</span><span class="n">release</span><span class="p">()</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">baton</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span>
    <span class="k">elif</span> <span class="n">verbose</span><span class="p">:</span>
        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;No modifications detected for re-loaded extension &#39;</span>
              <span class="s1">&#39;module </span><span class="si">{}</span><span class="s1">, skipping build step...&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">))</span>

    <span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Loading extension module </span><span class="si">{}</span><span class="s1">...&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">_import_module_from_library</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">build_directory</span><span class="p">,</span> <span class="n">is_python_module</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_write_ninja_file_and_compile_objects</span><span class="p">(</span>
        <span class="n">sources</span><span class="p">,</span>
        <span class="n">objects</span><span class="p">,</span>
        <span class="n">cflags</span><span class="p">,</span>
        <span class="n">post_cflags</span><span class="p">,</span>
        <span class="n">cuda_cflags</span><span class="p">,</span>
        <span class="n">cuda_post_cflags</span><span class="p">,</span>
        <span class="n">build_directory</span><span class="p">,</span>
        <span class="n">verbose</span><span class="p">,</span>
        <span class="n">with_cuda</span><span class="p">):</span>
    <span class="n">verify_ninja_availability</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
        <span class="n">compiler</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CXX&#39;</span><span class="p">,</span> <span class="s1">&#39;cl&#39;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">compiler</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CXX&#39;</span><span class="p">,</span> <span class="s1">&#39;c++&#39;</span><span class="p">)</span>
    <span class="n">check_compiler_abi_compatibility</span><span class="p">(</span><span class="n">compiler</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">with_cuda</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">with_cuda</span> <span class="o">=</span> <span class="nb">any</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="n">_is_cuda_file</span><span class="p">,</span> <span class="n">sources</span><span class="p">))</span>
    <span class="n">build_file_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">build_directory</span><span class="p">,</span> <span class="s1">&#39;build.ninja&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
        <span class="nb">print</span><span class="p">(</span>
            <span class="s1">&#39;Emitting ninja build file </span><span class="si">{}</span><span class="s1">...&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">build_file_path</span><span class="p">))</span>
    <span class="n">_write_ninja_file</span><span class="p">(</span>
        <span class="n">path</span><span class="o">=</span><span class="n">build_file_path</span><span class="p">,</span>
        <span class="n">cflags</span><span class="o">=</span><span class="n">cflags</span><span class="p">,</span>
        <span class="n">post_cflags</span><span class="o">=</span><span class="n">post_cflags</span><span class="p">,</span>
        <span class="n">cuda_cflags</span><span class="o">=</span><span class="n">cuda_cflags</span><span class="p">,</span>
        <span class="n">cuda_post_cflags</span><span class="o">=</span><span class="n">cuda_post_cflags</span><span class="p">,</span>
        <span class="n">sources</span><span class="o">=</span><span class="n">sources</span><span class="p">,</span>
        <span class="n">objects</span><span class="o">=</span><span class="n">objects</span><span class="p">,</span>
        <span class="n">ldflags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">library_target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">with_cuda</span><span class="o">=</span><span class="n">with_cuda</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Compiling objects...&#39;</span><span class="p">)</span>
    <span class="n">_run_ninja_build</span><span class="p">(</span>
        <span class="n">build_directory</span><span class="p">,</span>
        <span class="n">verbose</span><span class="p">,</span>
        <span class="c1"># It would be better if we could tell users the name of the extension</span>
        <span class="c1"># that failed to build but there isn&#39;t a good way to get it here.</span>
        <span class="n">error_prefix</span><span class="o">=</span><span class="s1">&#39;Error compiling objects for extension&#39;</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_write_ninja_file_and_build_library</span><span class="p">(</span>
        <span class="n">name</span><span class="p">,</span>
        <span class="n">sources</span><span class="p">,</span>
        <span class="n">extra_cflags</span><span class="p">,</span>
        <span class="n">extra_cuda_cflags</span><span class="p">,</span>
        <span class="n">extra_ldflags</span><span class="p">,</span>
        <span class="n">extra_include_paths</span><span class="p">,</span>
        <span class="n">build_directory</span><span class="p">,</span>
        <span class="n">verbose</span><span class="p">,</span>
        <span class="n">with_cuda</span><span class="p">):</span>
    <span class="n">verify_ninja_availability</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
        <span class="n">compiler</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CXX&#39;</span><span class="p">,</span> <span class="s1">&#39;cl&#39;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">compiler</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CXX&#39;</span><span class="p">,</span> <span class="s1">&#39;c++&#39;</span><span class="p">)</span>
    <span class="n">check_compiler_abi_compatibility</span><span class="p">(</span><span class="n">compiler</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">with_cuda</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">with_cuda</span> <span class="o">=</span> <span class="nb">any</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="n">_is_cuda_file</span><span class="p">,</span> <span class="n">sources</span><span class="p">))</span>
    <span class="n">extra_ldflags</span> <span class="o">=</span> <span class="n">_prepare_ldflags</span><span class="p">(</span>
        <span class="n">extra_ldflags</span> <span class="ow">or</span> <span class="p">[],</span>
        <span class="n">with_cuda</span><span class="p">,</span>
        <span class="n">verbose</span><span class="p">)</span>
    <span class="n">build_file_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">build_directory</span><span class="p">,</span> <span class="s1">&#39;build.ninja&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
        <span class="nb">print</span><span class="p">(</span>
            <span class="s1">&#39;Emitting ninja build file </span><span class="si">{}</span><span class="s1">...&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">build_file_path</span><span class="p">))</span>
    <span class="c1"># NOTE: Emitting a new ninja build file does not cause re-compilation if</span>
    <span class="c1"># the sources did not change, so it&#39;s ok to re-emit (and it&#39;s fast).</span>
    <span class="n">_write_ninja_file_to_build_library</span><span class="p">(</span>
        <span class="n">path</span><span class="o">=</span><span class="n">build_file_path</span><span class="p">,</span>
        <span class="n">name</span><span class="o">=</span><span class="n">name</span><span class="p">,</span>
        <span class="n">sources</span><span class="o">=</span><span class="n">sources</span><span class="p">,</span>
        <span class="n">extra_cflags</span><span class="o">=</span><span class="n">extra_cflags</span> <span class="ow">or</span> <span class="p">[],</span>
        <span class="n">extra_cuda_cflags</span><span class="o">=</span><span class="n">extra_cuda_cflags</span> <span class="ow">or</span> <span class="p">[],</span>
        <span class="n">extra_ldflags</span><span class="o">=</span><span class="n">extra_ldflags</span> <span class="ow">or</span> <span class="p">[],</span>
        <span class="n">extra_include_paths</span><span class="o">=</span><span class="n">extra_include_paths</span> <span class="ow">or</span> <span class="p">[],</span>
        <span class="n">with_cuda</span><span class="o">=</span><span class="n">with_cuda</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Building extension module </span><span class="si">{}</span><span class="s1">...&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">))</span>
    <span class="n">_run_ninja_build</span><span class="p">(</span>
        <span class="n">build_directory</span><span class="p">,</span>
        <span class="n">verbose</span><span class="p">,</span>
        <span class="n">error_prefix</span><span class="o">=</span><span class="s2">&quot;Error building extension &#39;</span><span class="si">{}</span><span class="s2">&#39;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">))</span>


<span class="k">def</span> <span class="nf">_is_ninja_available</span><span class="p">():</span>
    <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">devnull</span><span class="p">,</span> <span class="s1">&#39;wb&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">devnull</span><span class="p">:</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">subprocess</span><span class="o">.</span><span class="n">check_call</span><span class="p">(</span><span class="s1">&#39;ninja --version&#39;</span><span class="o">.</span><span class="n">split</span><span class="p">(),</span> <span class="n">stdout</span><span class="o">=</span><span class="n">devnull</span><span class="p">)</span>
        <span class="k">except</span> <span class="ne">OSError</span><span class="p">:</span>
            <span class="k">return</span> <span class="kc">False</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">return</span> <span class="kc">True</span>


<div class="viewcode-block" id="verify_ninja_availability"><a class="viewcode-back" href="../../../cpp_extension.html#torch.utils.cpp_extension.verify_ninja_availability">[docs]</a><span class="k">def</span> <span class="nf">verify_ninja_availability</span><span class="p">():</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Returns ``True`` if the `ninja &lt;https://ninja-build.org/&gt;`_ build system is</span>
<span class="sd">    available on the system.</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">_is_ninja_available</span><span class="p">():</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Ninja is required to load C++ extensions&quot;</span><span class="p">)</span></div>


<span class="k">def</span> <span class="nf">_prepare_ldflags</span><span class="p">(</span><span class="n">extra_ldflags</span><span class="p">,</span> <span class="n">with_cuda</span><span class="p">,</span> <span class="n">verbose</span><span class="p">):</span>
    <span class="n">here</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="vm">__file__</span><span class="p">)</span>
    <span class="n">torch_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">here</span><span class="p">))</span>
    <span class="n">lib_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">torch_path</span><span class="p">,</span> <span class="s1">&#39;lib&#39;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
        <span class="n">python_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">sys</span><span class="o">.</span><span class="n">executable</span><span class="p">)</span>
        <span class="n">python_lib_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">python_path</span><span class="p">,</span> <span class="s1">&#39;libs&#39;</span><span class="p">)</span>

        <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;c10.lib&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
            <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;c10_cuda.lib&#39;</span><span class="p">)</span>
        <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch_cpu.lib&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
            <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch_cuda.lib&#39;</span><span class="p">)</span>
            <span class="c1"># /INCLUDE is used to ensure torch_cuda is linked against in a project that relies on it.</span>
            <span class="c1"># Related issue: https://github.com/pytorch/pytorch/issues/31611</span>
            <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-INCLUDE:?warp_size@cuda@at@@YAHXZ&#39;</span><span class="p">)</span>
        <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch.lib&#39;</span><span class="p">)</span>
        <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;torch_python.lib&#39;</span><span class="p">)</span>
        <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;/LIBPATH:</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">python_lib_path</span><span class="p">))</span>
        <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;/LIBPATH:</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">lib_path</span><span class="p">))</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-L</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">lib_path</span><span class="p">))</span>
        <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-lc10&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
            <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-lc10_cuda&#39;</span><span class="p">)</span>
        <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-ltorch_cpu&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
            <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-ltorch_cuda&#39;</span><span class="p">)</span>
        <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-ltorch&#39;</span><span class="p">)</span>
        <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-ltorch_python&#39;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
            <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Detected CUDA files, patching ldflags&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
            <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;/LIBPATH:</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
                <span class="n">_join_cuda_home</span><span class="p">(</span><span class="s1">&#39;lib/x64&#39;</span><span class="p">)))</span>
            <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;cudart.lib&#39;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">CUDNN_HOME</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">CUDNN_HOME</span><span class="p">,</span> <span class="s1">&#39;lib/x64&#39;</span><span class="p">))</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-L</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">_join_cuda_home</span><span class="p">(</span><span class="s1">&#39;lib64&#39;</span><span class="p">)))</span>
            <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-lcudart&#39;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">CUDNN_HOME</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">extra_ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-L</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">CUDNN_HOME</span><span class="p">,</span> <span class="s1">&#39;lib64&#39;</span><span class="p">)))</span>

    <span class="k">return</span> <span class="n">extra_ldflags</span>


<span class="k">def</span> <span class="nf">_get_cuda_arch_flags</span><span class="p">(</span><span class="n">cflags</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Determine CUDA arch flags to use.</span>

<span class="sd">    For an arch, say &quot;6.1&quot;, the added compile flag will be</span>
<span class="sd">    ``-gencode=arch=compute_61,code=sm_61``.</span>
<span class="sd">    For an added &quot;+PTX&quot;, an additional</span>
<span class="sd">    ``-gencode=arch=compute_xx,code=compute_xx`` is added.</span>

<span class="sd">    See select_compute_arch.cmake for corresponding named and supported arches</span>
<span class="sd">    when building with CMake.</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="c1"># If cflags is given, there may already be user-provided arch flags in it</span>
    <span class="c1"># (from `extra_compile_args`)</span>
    <span class="k">if</span> <span class="n">cflags</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">flag</span> <span class="ow">in</span> <span class="n">cflags</span><span class="p">:</span>
            <span class="k">if</span> <span class="s1">&#39;arch&#39;</span> <span class="ow">in</span> <span class="n">flag</span><span class="p">:</span>
                <span class="k">return</span> <span class="p">[]</span>

    <span class="c1"># Note: keep combined names (&quot;arch1+arch2&quot;) above single names, otherwise</span>
    <span class="c1"># string replacement may not do the right thing</span>
    <span class="n">named_arches</span> <span class="o">=</span> <span class="n">collections</span><span class="o">.</span><span class="n">OrderedDict</span><span class="p">([</span>
        <span class="p">(</span><span class="s1">&#39;Kepler+Tesla&#39;</span><span class="p">,</span> <span class="s1">&#39;3.7&#39;</span><span class="p">),</span>
        <span class="p">(</span><span class="s1">&#39;Kepler&#39;</span><span class="p">,</span> <span class="s1">&#39;3.5+PTX&#39;</span><span class="p">),</span>
        <span class="p">(</span><span class="s1">&#39;Maxwell+Tegra&#39;</span><span class="p">,</span> <span class="s1">&#39;5.3&#39;</span><span class="p">),</span>
        <span class="p">(</span><span class="s1">&#39;Maxwell&#39;</span><span class="p">,</span> <span class="s1">&#39;5.0;5.2+PTX&#39;</span><span class="p">),</span>
        <span class="p">(</span><span class="s1">&#39;Pascal&#39;</span><span class="p">,</span> <span class="s1">&#39;6.0;6.1+PTX&#39;</span><span class="p">),</span>
        <span class="p">(</span><span class="s1">&#39;Volta&#39;</span><span class="p">,</span> <span class="s1">&#39;7.0+PTX&#39;</span><span class="p">),</span>
        <span class="p">(</span><span class="s1">&#39;Turing&#39;</span><span class="p">,</span> <span class="s1">&#39;7.5+PTX&#39;</span><span class="p">),</span>
    <span class="p">])</span>

    <span class="n">supported_arches</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;3.5&#39;</span><span class="p">,</span> <span class="s1">&#39;3.7&#39;</span><span class="p">,</span> <span class="s1">&#39;5.0&#39;</span><span class="p">,</span> <span class="s1">&#39;5.2&#39;</span><span class="p">,</span> <span class="s1">&#39;5.3&#39;</span><span class="p">,</span> <span class="s1">&#39;6.0&#39;</span><span class="p">,</span> <span class="s1">&#39;6.1&#39;</span><span class="p">,</span> <span class="s1">&#39;6.2&#39;</span><span class="p">,</span>
                        <span class="s1">&#39;7.0&#39;</span><span class="p">,</span> <span class="s1">&#39;7.2&#39;</span><span class="p">,</span> <span class="s1">&#39;7.5&#39;</span><span class="p">]</span>
    <span class="n">valid_arch_strings</span> <span class="o">=</span> <span class="n">supported_arches</span> <span class="o">+</span> <span class="p">[</span><span class="n">s</span> <span class="o">+</span> <span class="s2">&quot;+PTX&quot;</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">supported_arches</span><span class="p">]</span>

    <span class="c1"># The default is sm_30 for CUDA 9.x and 10.x</span>
    <span class="c1"># First check for an env var (same as used by the main setup.py)</span>
    <span class="c1"># Can be one or more architectures, e.g. &quot;6.1&quot; or &quot;3.5;5.2;6.0;6.1;7.0+PTX&quot;</span>
    <span class="c1"># See cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake</span>
    <span class="n">arch_list</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;TORCH_CUDA_ARCH_LIST&#39;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>

    <span class="c1"># If not given, determine what&#39;s needed for the GPU that can be found</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">arch_list</span><span class="p">:</span>
        <span class="n">capability</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">get_device_capability</span><span class="p">()</span>
        <span class="n">arch_list</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;</span><span class="si">{}</span><span class="s1">.</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">capability</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">capability</span><span class="p">[</span><span class="mi">1</span><span class="p">])]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="c1"># Deal with lists that are &#39; &#39; separated (only deal with &#39;;&#39; after)</span>
        <span class="n">arch_list</span> <span class="o">=</span> <span class="n">arch_list</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s1">&#39; &#39;</span><span class="p">,</span> <span class="s1">&#39;;&#39;</span><span class="p">)</span>
        <span class="c1"># Expand named arches</span>
        <span class="k">for</span> <span class="n">named_arch</span><span class="p">,</span> <span class="n">archval</span> <span class="ow">in</span> <span class="n">named_arches</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
            <span class="n">arch_list</span> <span class="o">=</span> <span class="n">arch_list</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">named_arch</span><span class="p">,</span> <span class="n">archval</span><span class="p">)</span>

        <span class="n">arch_list</span> <span class="o">=</span> <span class="n">arch_list</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s1">&#39;;&#39;</span><span class="p">)</span>

    <span class="n">flags</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">arch</span> <span class="ow">in</span> <span class="n">arch_list</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">arch</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">valid_arch_strings</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Unknown CUDA arch (</span><span class="si">{}</span><span class="s2">) or GPU not supported&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">arch</span><span class="p">))</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">num</span> <span class="o">=</span> <span class="n">arch</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="n">arch</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span>
            <span class="n">flags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-gencode=arch=compute_</span><span class="si">{}</span><span class="s1">,code=sm_</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">num</span><span class="p">,</span> <span class="n">num</span><span class="p">))</span>
            <span class="k">if</span> <span class="n">arch</span><span class="o">.</span><span class="n">endswith</span><span class="p">(</span><span class="s1">&#39;+PTX&#39;</span><span class="p">):</span>
                <span class="n">flags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-gencode=arch=compute_</span><span class="si">{}</span><span class="s1">,code=compute_</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">num</span><span class="p">,</span> <span class="n">num</span><span class="p">))</span>

    <span class="k">return</span> <span class="nb">list</span><span class="p">(</span><span class="nb">set</span><span class="p">(</span><span class="n">flags</span><span class="p">))</span>


<span class="k">def</span> <span class="nf">_get_rocm_arch_flags</span><span class="p">(</span><span class="n">cflags</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="c1"># If cflags is given, there may already be user-provided arch flags in it</span>
    <span class="c1"># (from `extra_compile_args`)</span>
    <span class="k">if</span> <span class="n">cflags</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">flag</span> <span class="ow">in</span> <span class="n">cflags</span><span class="p">:</span>
            <span class="k">if</span> <span class="s1">&#39;amdgpu-target&#39;</span> <span class="ow">in</span> <span class="n">flag</span><span class="p">:</span>
                <span class="k">return</span> <span class="p">[</span><span class="s1">&#39;-fno-gpu-rdc&#39;</span><span class="p">]</span>
    <span class="k">return</span> <span class="p">[</span>
        <span class="s1">&#39;--amdgpu-target=gfx803&#39;</span><span class="p">,</span>
        <span class="s1">&#39;--amdgpu-target=gfx900&#39;</span><span class="p">,</span>
        <span class="s1">&#39;--amdgpu-target=gfx906&#39;</span><span class="p">,</span>
        <span class="s1">&#39;--amdgpu-target=gfx908&#39;</span><span class="p">,</span>
        <span class="s1">&#39;-fno-gpu-rdc&#39;</span>
    <span class="p">]</span>


<span class="k">def</span> <span class="nf">_get_build_directory</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">verbose</span><span class="p">):</span>
    <span class="n">root_extensions_directory</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;TORCH_EXTENSIONS_DIR&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">root_extensions_directory</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">root_extensions_directory</span> <span class="o">=</span> <span class="n">get_default_build_root</span><span class="p">()</span>

    <span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Using </span><span class="si">{}</span><span class="s1"> as PyTorch extensions root...&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
            <span class="n">root_extensions_directory</span><span class="p">))</span>

    <span class="n">build_directory</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">root_extensions_directory</span><span class="p">,</span> <span class="n">name</span><span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">build_directory</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
            <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Creating extension directory </span><span class="si">{}</span><span class="s1">...&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">build_directory</span><span class="p">))</span>
        <span class="c1"># This is like mkdir -p, i.e. will also create parent directories.</span>
        <span class="n">os</span><span class="o">.</span><span class="n">makedirs</span><span class="p">(</span><span class="n">build_directory</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">build_directory</span>


<span class="k">def</span> <span class="nf">_get_num_workers</span><span class="p">(</span><span class="n">verbose</span><span class="p">):</span>
    <span class="n">max_jobs</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;MAX_JOBS&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">max_jobs</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">max_jobs</span><span class="o">.</span><span class="n">isdigit</span><span class="p">():</span>
        <span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
            <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Using envvar MAX_JOBS (</span><span class="si">{}</span><span class="s1">) as the number of workers...&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">max_jobs</span><span class="p">))</span>
        <span class="k">return</span> <span class="nb">int</span><span class="p">(</span><span class="n">max_jobs</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Allowing ninja to set a default number of workers... &#39;</span>
              <span class="s1">&#39;(overridable by setting the environment variable MAX_JOBS=N)&#39;</span><span class="p">)</span>
    <span class="k">return</span> <span class="kc">None</span>


<span class="k">def</span> <span class="nf">_run_ninja_build</span><span class="p">(</span><span class="n">build_directory</span><span class="p">,</span> <span class="n">verbose</span><span class="p">,</span> <span class="n">error_prefix</span><span class="p">):</span>
    <span class="n">command</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;ninja&#39;</span><span class="p">,</span> <span class="s1">&#39;-v&#39;</span><span class="p">]</span>
    <span class="n">num_workers</span> <span class="o">=</span> <span class="n">_get_num_workers</span><span class="p">(</span><span class="n">verbose</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">num_workers</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">command</span><span class="o">.</span><span class="n">extend</span><span class="p">([</span><span class="s1">&#39;-j&#39;</span><span class="p">,</span> <span class="nb">str</span><span class="p">(</span><span class="n">num_workers</span><span class="p">)])</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="n">sys</span><span class="o">.</span><span class="n">stdout</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span>
        <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">sys</span><span class="o">.</span><span class="n">version_info</span> <span class="o">&gt;=</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">5</span><span class="p">):</span>
            <span class="c1"># Warning: don&#39;t pass stdout=None to subprocess.run to get output.</span>
            <span class="c1"># subprocess.run assumes that sys.__stdout__ has not been modified and</span>
            <span class="c1"># attempts to write to it by default.  However, when we call _run_ninja_build</span>
            <span class="c1"># from ahead-of-time cpp extensions, the following happens:</span>
            <span class="c1"># 1) If the stdout encoding is not utf-8, setuptools detachs __stdout__.</span>
            <span class="c1">#    https://github.com/pypa/setuptools/blob/7e97def47723303fafabe48b22168bbc11bb4821/setuptools/dist.py#L1110</span>
            <span class="c1">#    (it probably shouldn&#39;t do this)</span>
            <span class="c1"># 2) subprocess.run (on POSIX, with no stdout override) relies on</span>
            <span class="c1">#    __stdout__ not being detached:</span>
            <span class="c1">#    https://github.com/python/cpython/blob/c352e6c7446c894b13643f538db312092b351789/Lib/subprocess.py#L1214</span>
            <span class="c1"># To work around this, we pass in the fileno directly and hope that</span>
            <span class="c1"># it is valid.</span>
            <span class="n">stdout_fileno</span> <span class="o">=</span> <span class="mi">1</span>
            <span class="n">subprocess</span><span class="o">.</span><span class="n">run</span><span class="p">(</span>
                <span class="n">command</span><span class="p">,</span>
                <span class="n">stdout</span><span class="o">=</span><span class="n">stdout_fileno</span> <span class="k">if</span> <span class="n">verbose</span> <span class="k">else</span> <span class="n">subprocess</span><span class="o">.</span><span class="n">PIPE</span><span class="p">,</span>
                <span class="n">stderr</span><span class="o">=</span><span class="n">subprocess</span><span class="o">.</span><span class="n">STDOUT</span><span class="p">,</span>
                <span class="n">cwd</span><span class="o">=</span><span class="n">build_directory</span><span class="p">,</span>
                <span class="n">check</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">subprocess</span><span class="o">.</span><span class="n">check_output</span><span class="p">(</span>
                <span class="n">command</span><span class="p">,</span>
                <span class="n">stderr</span><span class="o">=</span><span class="n">subprocess</span><span class="o">.</span><span class="n">STDOUT</span><span class="p">,</span>
                <span class="n">cwd</span><span class="o">=</span><span class="n">build_directory</span><span class="p">)</span>
    <span class="k">except</span> <span class="n">subprocess</span><span class="o">.</span><span class="n">CalledProcessError</span><span class="p">:</span>
        <span class="c1"># Python 2 and 3 compatible way of getting the error object.</span>
        <span class="n">_</span><span class="p">,</span> <span class="n">error</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">exc_info</span><span class="p">()</span>
        <span class="c1"># error.output contains the stdout and stderr of the build attempt.</span>
        <span class="n">message</span> <span class="o">=</span> <span class="n">error_prefix</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">error</span><span class="p">,</span> <span class="s1">&#39;output&#39;</span><span class="p">)</span> <span class="ow">and</span> <span class="n">error</span><span class="o">.</span><span class="n">output</span><span class="p">:</span>
            <span class="n">message</span> <span class="o">+=</span> <span class="s2">&quot;: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">error</span><span class="o">.</span><span class="n">output</span><span class="o">.</span><span class="n">decode</span><span class="p">())</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="n">message</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_import_module_from_library</span><span class="p">(</span><span class="n">module_name</span><span class="p">,</span> <span class="n">path</span><span class="p">,</span> <span class="n">is_python_module</span><span class="p">):</span>
    <span class="c1"># https://stackoverflow.com/questions/67631/how-to-import-a-module-given-the-full-path</span>
    <span class="n">file</span><span class="p">,</span> <span class="n">path</span><span class="p">,</span> <span class="n">description</span> <span class="o">=</span> <span class="n">imp</span><span class="o">.</span><span class="n">find_module</span><span class="p">(</span><span class="n">module_name</span><span class="p">,</span> <span class="p">[</span><span class="n">path</span><span class="p">])</span>
    <span class="c1"># Close the .so file after load.</span>
    <span class="k">with</span> <span class="n">file</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">is_python_module</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">imp</span><span class="o">.</span><span class="n">load_module</span><span class="p">(</span><span class="n">module_name</span><span class="p">,</span> <span class="n">file</span><span class="p">,</span> <span class="n">path</span><span class="p">,</span> <span class="n">description</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">load_library</span><span class="p">(</span><span class="n">path</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_write_ninja_file_to_build_library</span><span class="p">(</span><span class="n">path</span><span class="p">,</span>
                                       <span class="n">name</span><span class="p">,</span>
                                       <span class="n">sources</span><span class="p">,</span>
                                       <span class="n">extra_cflags</span><span class="p">,</span>
                                       <span class="n">extra_cuda_cflags</span><span class="p">,</span>
                                       <span class="n">extra_ldflags</span><span class="p">,</span>
                                       <span class="n">extra_include_paths</span><span class="p">,</span>
                                       <span class="n">with_cuda</span><span class="p">):</span>
    <span class="n">extra_cflags</span> <span class="o">=</span> <span class="p">[</span><span class="n">flag</span><span class="o">.</span><span class="n">strip</span><span class="p">()</span> <span class="k">for</span> <span class="n">flag</span> <span class="ow">in</span> <span class="n">extra_cflags</span><span class="p">]</span>
    <span class="n">extra_cuda_cflags</span> <span class="o">=</span> <span class="p">[</span><span class="n">flag</span><span class="o">.</span><span class="n">strip</span><span class="p">()</span> <span class="k">for</span> <span class="n">flag</span> <span class="ow">in</span> <span class="n">extra_cuda_cflags</span><span class="p">]</span>
    <span class="n">extra_ldflags</span> <span class="o">=</span> <span class="p">[</span><span class="n">flag</span><span class="o">.</span><span class="n">strip</span><span class="p">()</span> <span class="k">for</span> <span class="n">flag</span> <span class="ow">in</span> <span class="n">extra_ldflags</span><span class="p">]</span>
    <span class="n">extra_include_paths</span> <span class="o">=</span> <span class="p">[</span><span class="n">flag</span><span class="o">.</span><span class="n">strip</span><span class="p">()</span> <span class="k">for</span> <span class="n">flag</span> <span class="ow">in</span> <span class="n">extra_include_paths</span><span class="p">]</span>

    <span class="c1"># Turn into absolute paths so we can emit them into the ninja build</span>
    <span class="c1"># file wherever it is.</span>
    <span class="n">user_includes</span> <span class="o">=</span> <span class="p">[</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="n">file</span><span class="p">)</span> <span class="k">for</span> <span class="n">file</span> <span class="ow">in</span> <span class="n">extra_include_paths</span><span class="p">]</span>

    <span class="c1"># include_paths() gives us the location of torch/extension.h</span>
    <span class="n">system_includes</span> <span class="o">=</span> <span class="n">include_paths</span><span class="p">(</span><span class="n">with_cuda</span><span class="p">)</span>
    <span class="c1"># sysconfig.get_paths()[&#39;include&#39;] gives us the location of Python.h</span>
    <span class="n">system_includes</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">sysconfig</span><span class="o">.</span><span class="n">get_paths</span><span class="p">()[</span><span class="s1">&#39;include&#39;</span><span class="p">])</span>

    <span class="c1"># Windows does not understand `-isystem`.</span>
    <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
        <span class="n">user_includes</span> <span class="o">+=</span> <span class="n">system_includes</span>
        <span class="n">system_includes</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>

    <span class="n">common_cflags</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;-DTORCH_EXTENSION_NAME=</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">)]</span>
    <span class="n">common_cflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-DTORCH_API_INCLUDE_EXTENSION_H&#39;</span><span class="p">)</span>
    <span class="n">common_cflags</span> <span class="o">+=</span> <span class="p">[</span><span class="s1">&#39;-I</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">include</span><span class="p">)</span> <span class="k">for</span> <span class="n">include</span> <span class="ow">in</span> <span class="n">user_includes</span><span class="p">]</span>
    <span class="n">common_cflags</span> <span class="o">+=</span> <span class="p">[</span><span class="s1">&#39;-isystem </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">include</span><span class="p">)</span> <span class="k">for</span> <span class="n">include</span> <span class="ow">in</span> <span class="n">system_includes</span><span class="p">]</span>

    <span class="n">common_cflags</span> <span class="o">+=</span> <span class="p">[</span><span class="s1">&#39;-D_GLIBCXX_USE_CXX11_ABI=&#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_GLIBCXX_USE_CXX11_ABI</span><span class="p">))]</span>

    <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
        <span class="n">cflags</span> <span class="o">=</span> <span class="n">common_cflags</span> <span class="o">+</span> <span class="n">COMMON_MSVC_FLAGS</span> <span class="o">+</span> <span class="n">extra_cflags</span>
        <span class="kn">from</span> <span class="nn">distutils.spawn</span> <span class="kn">import</span> <span class="n">_nt_quote_args</span>
        <span class="n">cflags</span> <span class="o">=</span> <span class="n">_nt_quote_args</span><span class="p">(</span><span class="n">cflags</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">cflags</span> <span class="o">=</span> <span class="n">common_cflags</span> <span class="o">+</span> <span class="p">[</span><span class="s1">&#39;-fPIC&#39;</span><span class="p">,</span> <span class="s1">&#39;-std=c++14&#39;</span><span class="p">]</span> <span class="o">+</span> <span class="n">extra_cflags</span>

    <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
        <span class="n">cuda_flags</span> <span class="o">=</span> <span class="n">common_cflags</span> <span class="o">+</span> <span class="n">COMMON_NVCC_FLAGS</span> <span class="o">+</span> <span class="n">_get_cuda_arch_flags</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
            <span class="k">for</span> <span class="n">flag</span> <span class="ow">in</span> <span class="n">COMMON_MSVC_FLAGS</span><span class="p">:</span>
                <span class="n">cuda_flags</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;-Xcompiler&#39;</span><span class="p">,</span> <span class="n">flag</span><span class="p">]</span> <span class="o">+</span> <span class="n">cuda_flags</span>
            <span class="n">cuda_flags</span> <span class="o">=</span> <span class="n">_nt_quote_args</span><span class="p">(</span><span class="n">cuda_flags</span><span class="p">)</span>
            <span class="n">cuda_flags</span> <span class="o">+=</span> <span class="n">_nt_quote_args</span><span class="p">(</span><span class="n">extra_cuda_cflags</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">cuda_flags</span> <span class="o">+=</span> <span class="p">[</span><span class="s1">&#39;--compiler-options&#39;</span><span class="p">,</span> <span class="s2">&quot;&#39;-fPIC&#39;&quot;</span><span class="p">]</span>
            <span class="n">cuda_flags</span> <span class="o">+=</span> <span class="n">extra_cuda_cflags</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">flag</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s1">&#39;-std=&#39;</span><span class="p">)</span> <span class="k">for</span> <span class="n">flag</span> <span class="ow">in</span> <span class="n">cuda_flags</span><span class="p">):</span>
                <span class="n">cuda_flags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-std=c++14&#39;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">cuda_flags</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="nf">object_file_path</span><span class="p">(</span><span class="n">source_file</span><span class="p">):</span>
        <span class="c1"># &#39;/path/to/file.cpp&#39; -&gt; &#39;file&#39;</span>
        <span class="n">file_name</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">splitext</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">basename</span><span class="p">(</span><span class="n">source_file</span><span class="p">))[</span><span class="mi">0</span><span class="p">]</span>
        <span class="k">if</span> <span class="n">_is_cuda_file</span><span class="p">(</span><span class="n">source_file</span><span class="p">)</span> <span class="ow">and</span> <span class="n">with_cuda</span><span class="p">:</span>
            <span class="c1"># Use a different object filename in case a C++ and CUDA file have</span>
            <span class="c1"># the same filename but different extension (.cpp vs. .cu).</span>
            <span class="n">target</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="si">{}</span><span class="s1">.cuda.o&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">file_name</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">target</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="si">{}</span><span class="s1">.o&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">file_name</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">target</span>

    <span class="n">objects</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="n">object_file_path</span><span class="p">,</span> <span class="n">sources</span><span class="p">))</span>

    <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
        <span class="n">ldflags</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;/DLL&#39;</span><span class="p">]</span> <span class="o">+</span> <span class="n">extra_ldflags</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">ldflags</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;-shared&#39;</span><span class="p">]</span> <span class="o">+</span> <span class="n">extra_ldflags</span>
    <span class="c1"># The darwin linker needs explicit consent to ignore unresolved symbols.</span>
    <span class="k">if</span> <span class="n">sys</span><span class="o">.</span><span class="n">platform</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s1">&#39;darwin&#39;</span><span class="p">):</span>
        <span class="n">ldflags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;-undefined dynamic_lookup&#39;</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
        <span class="n">ldflags</span> <span class="o">=</span> <span class="n">_nt_quote_args</span><span class="p">(</span><span class="n">ldflags</span><span class="p">)</span>

    <span class="n">ext</span> <span class="o">=</span> <span class="s1">&#39;pyd&#39;</span> <span class="k">if</span> <span class="n">IS_WINDOWS</span> <span class="k">else</span> <span class="s1">&#39;so&#39;</span>
    <span class="n">library_target</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="si">{}</span><span class="s1">.</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">ext</span><span class="p">)</span>

    <span class="n">_write_ninja_file</span><span class="p">(</span>
        <span class="n">path</span><span class="o">=</span><span class="n">path</span><span class="p">,</span>
        <span class="n">cflags</span><span class="o">=</span><span class="n">cflags</span><span class="p">,</span>
        <span class="n">post_cflags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">cuda_cflags</span><span class="o">=</span><span class="n">cuda_flags</span><span class="p">,</span>
        <span class="n">cuda_post_cflags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">sources</span><span class="o">=</span><span class="n">sources</span><span class="p">,</span>
        <span class="n">objects</span><span class="o">=</span><span class="n">objects</span><span class="p">,</span>
        <span class="n">ldflags</span><span class="o">=</span><span class="n">ldflags</span><span class="p">,</span>
        <span class="n">library_target</span><span class="o">=</span><span class="n">library_target</span><span class="p">,</span>
        <span class="n">with_cuda</span><span class="o">=</span><span class="n">with_cuda</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_write_ninja_file</span><span class="p">(</span><span class="n">path</span><span class="p">,</span>
                      <span class="n">cflags</span><span class="p">,</span>
                      <span class="n">post_cflags</span><span class="p">,</span>
                      <span class="n">cuda_cflags</span><span class="p">,</span>
                      <span class="n">cuda_post_cflags</span><span class="p">,</span>
                      <span class="n">sources</span><span class="p">,</span>
                      <span class="n">objects</span><span class="p">,</span>
                      <span class="n">ldflags</span><span class="p">,</span>
                      <span class="n">library_target</span><span class="p">,</span>
                      <span class="n">with_cuda</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Write a ninja file that does the desired compiling and linking.</span>

<span class="sd">    `path`: Where to write this file</span>
<span class="sd">    `cflags`: list of flags to pass to $cxx. Can be None.</span>
<span class="sd">    `post_cflags`: list of flags to append to the $cxx invocation. Can be None.</span>
<span class="sd">    `cuda_cflags`: list of flags to pass to $nvcc. Can be None.</span>
<span class="sd">    `cuda_postflags`: list of flags to append to the $nvcc invocation. Can be None.</span>
<span class="sd">    `sources`: list of paths to source files</span>
<span class="sd">    `objects`: list of desired paths to objects, one per source.</span>
<span class="sd">    `ldflags`: list of flags to pass to linker. Can be None.</span>
<span class="sd">    `library_target`: Name of the output library. Can be None; in that case,</span>
<span class="sd">                      we do no linking.</span>
<span class="sd">    `with_cuda`: If we should be compiling with CUDA.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="nf">sanitize_flags</span><span class="p">(</span><span class="n">flags</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">flags</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">return</span> <span class="p">[]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">return</span> <span class="p">[</span><span class="n">flag</span><span class="o">.</span><span class="n">strip</span><span class="p">()</span> <span class="k">for</span> <span class="n">flag</span> <span class="ow">in</span> <span class="n">flags</span><span class="p">]</span>

    <span class="n">cflags</span> <span class="o">=</span> <span class="n">sanitize_flags</span><span class="p">(</span><span class="n">cflags</span><span class="p">)</span>
    <span class="n">post_cflags</span> <span class="o">=</span> <span class="n">sanitize_flags</span><span class="p">(</span><span class="n">post_cflags</span><span class="p">)</span>
    <span class="n">cuda_cflags</span> <span class="o">=</span> <span class="n">sanitize_flags</span><span class="p">(</span><span class="n">cuda_cflags</span><span class="p">)</span>
    <span class="n">cuda_post_cflags</span> <span class="o">=</span> <span class="n">sanitize_flags</span><span class="p">(</span><span class="n">cuda_post_cflags</span><span class="p">)</span>
    <span class="n">ldflags</span> <span class="o">=</span> <span class="n">sanitize_flags</span><span class="p">(</span><span class="n">ldflags</span><span class="p">)</span>

    <span class="c1"># Sanity checks...</span>
    <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">sources</span><span class="p">)</span> <span class="o">==</span> <span class="nb">len</span><span class="p">(</span><span class="n">objects</span><span class="p">)</span>
    <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">sources</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span>

    <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
        <span class="n">compiler</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CXX&#39;</span><span class="p">,</span> <span class="s1">&#39;cl&#39;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">compiler</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;CXX&#39;</span><span class="p">,</span> <span class="s1">&#39;c++&#39;</span><span class="p">)</span>

    <span class="c1"># Version 1.3 is required for the `deps` directive.</span>
    <span class="n">config</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;ninja_required_version = 1.3&#39;</span><span class="p">]</span>
    <span class="n">config</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;cxx = </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">compiler</span><span class="p">))</span>
    <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
        <span class="n">config</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;nvcc = </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">_join_cuda_home</span><span class="p">(</span><span class="s1">&#39;bin&#39;</span><span class="p">,</span> <span class="s1">&#39;nvcc&#39;</span><span class="p">)))</span>

    <span class="n">flags</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;cflags = </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="s1">&#39; &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">cflags</span><span class="p">))]</span>
    <span class="n">flags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;post_cflags = </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="s1">&#39; &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">post_cflags</span><span class="p">)))</span>
    <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
        <span class="n">flags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;cuda_cflags = </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="s1">&#39; &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">cuda_cflags</span><span class="p">)))</span>
        <span class="n">flags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;cuda_post_cflags = </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="s1">&#39; &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">cuda_post_cflags</span><span class="p">)))</span>
    <span class="n">flags</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;ldflags = </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="s1">&#39; &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">ldflags</span><span class="p">)))</span>

    <span class="c1"># Turn into absolute paths so we can emit them into the ninja build</span>
    <span class="c1"># file wherever it is.</span>
    <span class="n">sources</span> <span class="o">=</span> <span class="p">[</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="n">file</span><span class="p">)</span> <span class="k">for</span> <span class="n">file</span> <span class="ow">in</span> <span class="n">sources</span><span class="p">]</span>

    <span class="c1"># See https://ninja-build.org/build.ninja.html for reference.</span>
    <span class="n">compile_rule</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;rule compile&#39;</span><span class="p">]</span>
    <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
        <span class="n">compile_rule</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
            <span class="s1">&#39;  command = cl /showIncludes $cflags -c $in /Fo$out $post_cflags&#39;</span><span class="p">)</span>
        <span class="n">compile_rule</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;  deps = msvc&#39;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">compile_rule</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
            <span class="s1">&#39;  command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags&#39;</span><span class="p">)</span>
        <span class="n">compile_rule</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;  depfile = $out.d&#39;</span><span class="p">)</span>
        <span class="n">compile_rule</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;  deps = gcc&#39;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
        <span class="n">cuda_compile_rule</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;rule cuda_compile&#39;</span><span class="p">]</span>
        <span class="n">cuda_compile_rule</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
            <span class="s1">&#39;  command = $nvcc $cuda_cflags -c $in -o $out $cuda_post_cflags&#39;</span><span class="p">)</span>

    <span class="c1"># Emit one build rule per source to enable incremental build.</span>
    <span class="n">build</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">source_file</span><span class="p">,</span> <span class="n">object_file</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">sources</span><span class="p">,</span> <span class="n">objects</span><span class="p">):</span>
        <span class="n">is_cuda_source</span> <span class="o">=</span> <span class="n">_is_cuda_file</span><span class="p">(</span><span class="n">source_file</span><span class="p">)</span> <span class="ow">and</span> <span class="n">with_cuda</span>
        <span class="n">rule</span> <span class="o">=</span> <span class="s1">&#39;cuda_compile&#39;</span> <span class="k">if</span> <span class="n">is_cuda_source</span> <span class="k">else</span> <span class="s1">&#39;compile&#39;</span>
        <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
            <span class="n">source_file</span> <span class="o">=</span> <span class="n">source_file</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s1">&#39;:&#39;</span><span class="p">,</span> <span class="s1">&#39;$:&#39;</span><span class="p">)</span>
            <span class="n">object_file</span> <span class="o">=</span> <span class="n">object_file</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s1">&#39;:&#39;</span><span class="p">,</span> <span class="s1">&#39;$:&#39;</span><span class="p">)</span>
        <span class="n">source_file</span> <span class="o">=</span> <span class="n">source_file</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot; &quot;</span><span class="p">,</span> <span class="s2">&quot;$ &quot;</span><span class="p">)</span>
        <span class="n">build</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;build </span><span class="si">{}</span><span class="s1">: </span><span class="si">{}</span><span class="s1"> </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">object_file</span><span class="p">,</span> <span class="n">rule</span><span class="p">,</span> <span class="n">source_file</span><span class="p">))</span>

    <span class="k">if</span> <span class="n">library_target</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">link_rule</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;rule link&#39;</span><span class="p">]</span>
        <span class="k">if</span> <span class="n">IS_WINDOWS</span><span class="p">:</span>
            <span class="n">cl_paths</span> <span class="o">=</span> <span class="n">subprocess</span><span class="o">.</span><span class="n">check_output</span><span class="p">([</span><span class="s1">&#39;where&#39;</span><span class="p">,</span>
                                                <span class="s1">&#39;cl&#39;</span><span class="p">])</span><span class="o">.</span><span class="n">decode</span><span class="p">()</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s1">&#39;</span><span class="se">\r\n</span><span class="s1">&#39;</span><span class="p">)</span>
            <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">cl_paths</span><span class="p">)</span> <span class="o">&gt;=</span> <span class="mi">1</span><span class="p">:</span>
                <span class="n">cl_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">cl_paths</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s1">&#39;:&#39;</span><span class="p">,</span> <span class="s1">&#39;$:&#39;</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;MSVC is required to load C++ extensions&quot;</span><span class="p">)</span>
            <span class="n">link_rule</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
                <span class="s1">&#39;  command = &quot;</span><span class="si">{}</span><span class="s1">/link.exe&quot; $in /nologo $ldflags /out:$out&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
                    <span class="n">cl_path</span><span class="p">))</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">link_rule</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;  command = $cxx $in $ldflags -o $out&#39;</span><span class="p">)</span>

        <span class="n">link</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;build </span><span class="si">{}</span><span class="s1">: link </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">library_target</span><span class="p">,</span> <span class="s1">&#39; &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">objects</span><span class="p">))]</span>

        <span class="n">default</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;default </span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">library_target</span><span class="p">)]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">link_rule</span><span class="p">,</span> <span class="n">link</span><span class="p">,</span> <span class="n">default</span> <span class="o">=</span> <span class="p">[],</span> <span class="p">[],</span> <span class="p">[]</span>

    <span class="c1"># &#39;Blocks&#39; should be separated by newlines, for visual benefit.</span>
    <span class="n">blocks</span> <span class="o">=</span> <span class="p">[</span><span class="n">config</span><span class="p">,</span> <span class="n">flags</span><span class="p">,</span> <span class="n">compile_rule</span><span class="p">]</span>
    <span class="k">if</span> <span class="n">with_cuda</span><span class="p">:</span>
        <span class="n">blocks</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">cuda_compile_rule</span><span class="p">)</span>
    <span class="n">blocks</span> <span class="o">+=</span> <span class="p">[</span><span class="n">link_rule</span><span class="p">,</span> <span class="n">build</span><span class="p">,</span> <span class="n">link</span><span class="p">,</span> <span class="n">default</span><span class="p">]</span>
    <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="s1">&#39;w&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">build_file</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">block</span> <span class="ow">in</span> <span class="n">blocks</span><span class="p">:</span>
            <span class="n">lines</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="se">\n</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">block</span><span class="p">)</span>
            <span class="n">build_file</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s1">&#39;</span><span class="si">{}</span><span class="se">\n\n</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">lines</span><span class="p">))</span>


<span class="k">def</span> <span class="nf">_join_cuda_home</span><span class="p">(</span><span class="o">*</span><span class="n">paths</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Joins paths with CUDA_HOME, or raises an error if it CUDA_HOME is not set.</span>

<span class="sd">    This is basically a lazy way of raising an error for missing $CUDA_HOME</span>
<span class="sd">    only once we need to get any CUDA-specific path.</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="k">if</span> <span class="n">CUDA_HOME</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">EnvironmentError</span><span class="p">(</span><span class="s1">&#39;CUDA_HOME environment variable is not set. &#39;</span>
                               <span class="s1">&#39;Please set it to your CUDA install root.&#39;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">CUDA_HOME</span><span class="p">,</span> <span class="o">*</span><span class="n">paths</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_is_cuda_file</span><span class="p">(</span><span class="n">path</span><span class="p">):</span>
    <span class="n">valid_ext</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;.cu&#39;</span><span class="p">,</span> <span class="s1">&#39;.cuh&#39;</span><span class="p">]</span>
    <span class="k">if</span> <span class="n">IS_HIP_EXTENSION</span><span class="p">:</span>
        <span class="n">valid_ext</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;.hip&#39;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">splitext</span><span class="p">(</span><span class="n">path</span><span class="p">)[</span><span class="mi">1</span><span class="p">]</span> <span class="ow">in</span> <span class="n">valid_ext</span>
</pre></div>

             </article>
             
            </div>
            <footer>
  

  

    <hr>

  

  <div role="contentinfo">
    <p>
        &copy; Copyright 2019, Torch Contributors.

    </p>
  </div>
    
      <div>
        Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a> provided by <a href="https://readthedocs.org">Read the Docs</a>.
      </div>
     

</footer>

          </div>
        </div>

        <div class="pytorch-content-right" id="pytorch-content-right">
          <div class="pytorch-right-menu" id="pytorch-right-menu">
            <div class="pytorch-side-scroll" id="pytorch-side-scroll-right">
              
            </div>
          </div>
        </div>
      </section>
    </div>

  


  

     
       <script type="text/javascript" id="documentation_options" data-url_root="../../../" src="../../../_static/documentation_options.js"></script>
         <script src="../../../_static/jquery.js"></script>
         <script src="../../../_static/underscore.js"></script>
         <script src="../../../_static/doctools.js"></script>
         <script src="../../../_static/language_data.js"></script>
     

  

  <script type="text/javascript" src="../../../_static/js/vendor/popper.min.js"></script>
  <script type="text/javascript" src="../../../_static/js/vendor/bootstrap.min.js"></script>
  <script src="https://cdnjs.cloudflare.com/ajax/libs/list.js/1.5.0/list.min.js"></script>
  <script type="text/javascript" src="../../../_static/js/theme.js"></script>

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>
 
<script>
  (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){
  (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o),
  m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m)
  })(window,document,'script','https://www.google-analytics.com/analytics.js','ga');

  ga('create', 'UA-90545585-1', 'auto');
  ga('send', 'pageview');

</script>

<script async src="https://www.googletagmanager.com/gtag/js?id=UA-117752657-2"></script>

<script>
  window.dataLayer = window.dataLayer || [];

  function gtag(){dataLayer.push(arguments);}

  gtag('js', new Date());
  gtag('config', 'UA-117752657-2');
</script>

<img height="1" width="1" style="border-style:none;" alt="" src="https://www.googleadservices.com/pagead/conversion/795629140/?label=txkmCPmdtosBENSssfsC&amp;guid=ON&amp;script=0"/>


  <!-- Begin Footer -->

  <div class="container-fluid docs-tutorials-resources" id="docs-tutorials-resources">
    <div class="container">
      <div class="row">
        <div class="col-md-4 text-center">
          <h2>Docs</h2>
          <p>Access comprehensive developer documentation for PyTorch</p>
          <a class="with-right-arrow" href="https://pytorch.org/docs/stable/index.html">View Docs</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Tutorials</h2>
          <p>Get in-depth tutorials for beginners and advanced developers</p>
          <a class="with-right-arrow" href="https://pytorch.org/tutorials">View Tutorials</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Resources</h2>
          <p>Find development resources and get your questions answered</p>
          <a class="with-right-arrow" href="https://pytorch.org/resources">View Resources</a>
        </div>
      </div>
    </div>
  </div>

  <footer class="site-footer">
    <div class="container footer-container">
      <div class="footer-logo-wrapper">
        <a href="https://pytorch.org/" class="footer-logo"></a>
      </div>

      <div class="footer-links-wrapper">
        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/">PyTorch</a></li>
            <li><a href="https://pytorch.org/get-started">Get Started</a></li>
            <li><a href="https://pytorch.org/features">Features</a></li>
            <li><a href="https://pytorch.org/ecosystem">Ecosystem</a></li>
            <li><a href="https://pytorch.org/blog/">Blog</a></li>
            <li><a href="https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md">Contributing</a></li>
          </ul>
        </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/resources">Resources</a></li>
            <li><a href="https://pytorch.org/tutorials">Tutorials</a></li>
            <li><a href="https://pytorch.org/docs/stable/index.html">Docs</a></li>
            <li><a href="https://discuss.pytorch.org" target="_blank">Discuss</a></li>
            <li><a href="https://github.com/pytorch/pytorch/issues" target="_blank">Github Issues</a></li>
            <li><a href="https://pytorch.org/assets/brand-guidelines/PyTorch-Brand-Guidelines.pdf" target="_blank">Brand Guidelines</a></li>
          </ul>
        </div>

        <div class="footer-links-col follow-us-col">
          <ul>
            <li class="list-title">Stay Connected</li>
            <li>
              <div id="mc_embed_signup">
                <form
                  action="https://twitter.us14.list-manage.com/subscribe/post?u=75419c71fe0a935e53dfa4a3f&id=91d0dccd39"
                  method="post"
                  id="mc-embedded-subscribe-form"
                  name="mc-embedded-subscribe-form"
                  class="email-subscribe-form validate"
                  target="_blank"
                  novalidate>
                  <div id="mc_embed_signup_scroll" class="email-subscribe-form-fields-wrapper">
                    <div class="mc-field-group">
                      <label for="mce-EMAIL" style="display:none;">Email Address</label>
                      <input type="email" value="" name="EMAIL" class="required email" id="mce-EMAIL" placeholder="Email Address">
                    </div>

                    <div id="mce-responses" class="clear">
                      <div class="response" id="mce-error-response" style="display:none"></div>
                      <div class="response" id="mce-success-response" style="display:none"></div>
                    </div>    <!-- real people should not fill this in and expect good things - do not remove this or risk form bot signups-->

                    <div style="position: absolute; left: -5000px;" aria-hidden="true"><input type="text" name="b_75419c71fe0a935e53dfa4a3f_91d0dccd39" tabindex="-1" value=""></div>

                    <div class="clear">
                      <input type="submit" value="" name="subscribe" id="mc-embedded-subscribe" class="button email-subscribe-button">
                    </div>
                  </div>
                </form>
              </div>

            </li>
          </ul>

          <div class="footer-social-icons">
            <a href="https://www.facebook.com/pytorch" target="_blank" class="facebook"></a>
            <a href="https://twitter.com/pytorch" target="_blank" class="twitter"></a>
            <a href="https://www.youtube.com/pytorch" target="_blank" class="youtube"></a>
          </div>
        </div>
      </div>
    </div>
  </footer>

  <div class="cookie-banner-wrapper">
  <div class="container">
    <p class="gdpr-notice">To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: <a href="https://www.facebook.com/policies/cookies/">Cookies Policy</a>.</p>
    <img class="close-button" src="../../../_static/images/pytorch-x.svg">
  </div>
</div>

  <!-- End Footer -->

  <!-- Begin Mobile Menu -->

  <div class="mobile-main-menu">
    <div class="container-fluid">
      <div class="container">
        <div class="mobile-main-menu-header-container">
          <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>
          <a class="main-menu-close-button" href="#" data-behavior="close-mobile-menu"></a>
        </div>
      </div>
    </div>

    <div class="mobile-main-menu-links-container">
      <div class="main-menu">
        <ul>
          <li>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

          <li>
            <a href="https://pytorch.org/features">Features</a>
          </li>

          <li>
            <a href="https://pytorch.org/ecosystem">Ecosystem</a>
          </li>

          <li>
            <a href="https://pytorch.org/mobile">Mobile</a>
          </li>

          <li>
            <a href="https://pytorch.org/hub">PyTorch Hub</a>
          </li>

          <li>
            <a href="https://pytorch.org/blog/">Blog</a>
          </li>

          <li>
            <a href="https://pytorch.org/tutorials">Tutorials</a>
          </li>

          <li class="active">
            <a href="https://pytorch.org/docs/stable/index.html">Docs</a>
          </li>

          <li>
            <a href="https://pytorch.org/resources">Resources</a>
          </li>

          <li>
            <a href="https://github.com/pytorch/pytorch">Github</a>
          </li>
        </ul>
      </div>
    </div>
  </div>

  <!-- End Mobile Menu -->

  <script type="text/javascript" src="../../../_static/js/vendor/anchor.min.js"></script>

  <script type="text/javascript">
    $(document).ready(function() {
      mobileMenu.bind();
      mobileTOC.bind();
      pytorchAnchors.bind();
      sideMenus.bind();
      scrollToAnchor.bind();
      highlightNavigation.bind();
      mainMenuDropdown.bind();
      filterTags.bind();

      // Remove any empty p tags that Sphinx adds
      $("[data-tags='null']").remove();

      // Add class to links that have code blocks, since we cannot create links in code blocks
      $("article.pytorch-article a span.pre").each(function(e) {
        $(this).closest("a").addClass("has-code");
      });
    })
  </script>
</body>
</html>